better

Author: James Cor

server/handler.go

@@ -496,7 +496,7 @@ func (h *Handler) doQuery(
 	} else if analyzer.FlagIsSet(qFlags, sql.QFlagMax1Row) {
 		r, err = resultForMax1RowIter(sqlCtx, schema, rowIter, resultFields, buf)
 	} else if ri2, ok := rowIter.(sql.RowIter2); ok && ri2.IsRowIter2(sqlCtx) {
-		r, err = h.resultForDefaultIter2(sqlCtx, ri2, resultFields, callback, more)
+		r, processedAtLeastOneBatch, err = h.resultForDefaultIter2(sqlCtx, c, ri2, resultFields, callback, more)
 	} else {
 		r, processedAtLeastOneBatch, err = h.resultForDefaultIter(sqlCtx, c, schema, rowIter, callback, resultFields, more, buf)
 	}
@@ -770,30 +770,107 @@ func (h *Handler) resultForDefaultIter(ctx *sql.Context, c *mysql.Conn, schema s
 	return r, processedAtLeastOneBatch, nil
 }
 
-func (h *Handler) resultForDefaultIter2(ctx *sql.Context, iter sql.RowIter2, resultFields []*querypb.Field, callback func(*sqltypes.Result, bool) error, more bool) (*sqltypes.Result, error) {
-	res := &sqltypes.Result{Fields: resultFields}
-	for {
-		if res.RowsAffected == rowsBatch {
-			if err := callback(res, more); err != nil {
-				return nil, err
-			}
-			res = nil
-		}
-		row, err := iter.Next2(ctx)
-		if err == io.EOF {
-			return res, nil
+func (h *Handler) resultForDefaultIter2(ctx *sql.Context, c *mysql.Conn, iter sql.RowIter2, resultFields []*querypb.Field, callback func(*sqltypes.Result, bool) error, more bool) (*sqltypes.Result, bool, error) {
+	defer trace.StartRegion(ctx, "Handler.resultForDefaultIter").End()
+
+	eg, ctx := ctx.NewErrgroup()
+	pan2err := func(err *error) {
+		if recoveredPanic := recover(); recoveredPanic != nil {
+			stack := debug.Stack()
+			wrappedErr := fmt.Errorf("handler caught panic: %v\n%s", recoveredPanic, stack)
+			*err = goerrors.Join(*err, wrappedErr)
 		}
-		if err != nil {
-			return nil, err
+	}
+
+	// TODO: poll for closed connections should obviously also run even if
+	// we're doing something with an OK result or a single row result, etc.
+	// This should be in the caller.
+	pollCtx, cancelF := ctx.NewSubContext()
+	eg.Go(func() (err error) {
+		defer pan2err(&err)
+		return h.pollForClosedConnection(pollCtx, c)
+	})
+
+	// Default waitTime is one minute if there is no timeout configured, in which case
+	// it will loop to iterate again unless the socket died by the OS timeout or other problems.
+	// If there is a timeout, it will be enforced to ensure that Vitess has a chance to
+	// call Handler.CloseConnection()
+	waitTime := 1 * time.Minute
+	if h.readTimeout > 0 {
+		waitTime = h.readTimeout
+	}
+	timer := time.NewTimer(waitTime)
+	defer timer.Stop()
+
+	wg := sync.WaitGroup{}
+	wg.Add(1)
+
+	var res *sqltypes.Result
+	var processedAtLeastOneBatch bool
+	eg.Go(func() (err error) {
+		defer pan2err(&err)
+		defer cancelF()
+		defer wg.Done()
+		for {
+			if res == nil {
+				res = &sqltypes.Result{Fields: resultFields}
+			}
+			if res.RowsAffected == rowsBatch {
+				if err := callback(res, more); err != nil {
+					return err
+				}
+				res = nil
+				processedAtLeastOneBatch = true
+				continue
+			}
+
+			select {
+			case <-ctx.Done():
+				return context.Cause(ctx)
+			case <-timer.C:
+				// TODO: timer should probably go in its own thread, as rowChan is blocking
+				if h.readTimeout != 0 {
+					// Cancel and return so Vitess can call the CloseConnection callback
+					ctx.GetLogger().Tracef("connection timeout")
+					return ErrRowTimeout.New()
+				}
+			default:
+				row, err := iter.Next2(ctx)
+				if err == io.EOF {
+					return nil
+				}
+				if err != nil {
+					return err
+				}
+				outRow := make([]sqltypes.Value, len(row))
+				for i := range row {
+					outRow[i] = sqltypes.MakeTrusted(row[i].Typ, row[i].Val)
+				}
+				res.Rows = append(res.Rows, outRow)
+				res.RowsAffected++
+			}
+			timer.Reset(waitTime)
 		}
+	})
 
-		outRow := make([]sqltypes.Value, len(res.Rows))
-		for i := range row {
-			outRow[i] = sqltypes.MakeTrusted(row[i].Typ, row[i].Val)
+	// Close() kills this PID in the process list,
+	// wait until all rows have be sent over the wire
+	eg.Go(func() (err error) {
+		defer pan2err(&err)
+		wg.Wait()
+		return iter.Close(ctx)
+	})
+
+	err := eg.Wait()
+	if err != nil {
+		ctx.GetLogger().WithError(err).Warn("error running query")
+		if verboseErrorLogging {
+			fmt.Printf("Err: %+v", err)
 		}
-		res.Rows = append(res.Rows, outRow)
-		res.RowsAffected++
+		return nil, false, err
 	}
+
+	return res, processedAtLeastOneBatch, nil
 }
 
 // See https://dev.mysql.com/doc/internals/en/status-flags.html

sql/plan/filter.go

@@ -15,6 +15,7 @@
 package plan
 
 import (
+	"fmt"
 	"github.com/dolthub/go-mysql-server/sql"
 )
 
@@ -133,6 +134,29 @@ func (i *FilterIter) Next(ctx *sql.Context) (sql.Row, error) {
 	}
 }
 
+func (i *FilterIter) Next2(ctx *sql.Context) (sql.Row2, error) {
+	ri2, ok := i.childIter.(sql.RowIter2)
+	if !ok {
+		panic(fmt.Sprintf("%T is not a sql.RowIter2", i.childIter))
+	}
+
+	for {
+		row, err := ri2.Next(ctx)
+		if err != nil {
+			return nil, err
+		}
+
+		res, err := sql.EvaluateCondition(ctx, i.cond, row)
+		if err != nil {
+			return nil, err
+		}
+
+		if sql.IsTrue(res) {
+			return nil, nil
+		}
+	}
+}
+
 // Close implements the RowIter interface.
 func (i *FilterIter) Close(ctx *sql.Context) error {
 	return i.childIter.Close(ctx)