plan: compute all inner joins in memory if they fit

sql/plan/innerjoin.go

@@ -4,15 +4,52 @@ import (
 	"io"
 	"os"
 	"reflect"
+	"runtime"
+	"strconv"
+	"strings"
 
 	opentracing "github.com/opentracing/opentracing-go"
+	"github.com/pbnjay/memory"
+	"github.com/sirupsen/logrus"
 	"gopkg.in/src-d/go-mysql-server.v0/sql"
 )
 
-const experimentalInMemoryJoinKey = "EXPERIMENTAL_IN_MEMORY_JOIN"
-const inMemoryJoinSessionVar = "inmemory_joins"
+const (
+	experimentalInMemoryJoinKey = "EXPERIMENTAL_IN_MEMORY_JOIN"
+	maxMemoryJoinKey            = "MAX_MEMORY_INNER_JOIN"
+	inMemoryJoinSessionVar      = "inmemory_joins"
+	memoryThresholdSessionVar   = "max_memory_joins"
+)
+
+var (
+	useInMemoryJoins = shouldUseMemoryJoinsByEnv()
+	// One fifth of the total physical memory available on the OS (ignoring the
+	// memory used by other processes).
+	defaultMemoryThreshold = memory.TotalMemory() / 5
+	// Maximum amount of memory the gitbase server can have in use before
+	// considering all inner joins should be done using multipass mode.
+	maxMemoryJoin = loadMemoryThreshold()
+)
+
+func shouldUseMemoryJoinsByEnv() bool {
+	v := strings.TrimSpace(strings.ToLower(os.Getenv(experimentalInMemoryJoinKey)))
+	return v == "on" || v == "1"
+}
+
+func loadMemoryThreshold() uint64 {
+	v, ok := os.LookupEnv(maxMemoryJoinKey)
+	if !ok {
+		return defaultMemoryThreshold
+	}
 
-var useInMemoryJoins = os.Getenv(experimentalInMemoryJoinKey) != ""
+	n, err := strconv.ParseUint(v, 10, 64)
+	if err != nil {
+		logrus.Warnf("invalid value %q given to %s environment variable", v, maxMemoryJoinKey)
+		return defaultMemoryThreshold
+	}
+
+	return n
+}
 
 // InnerJoin is an inner join between two tables.
 type InnerJoin struct {
@@ -73,27 +110,17 @@ func (j *InnerJoin) RowIter(ctx *sql.Context) (sql.RowIter, error) {
 		inMemorySession = true
 	}
 
-	var iter sql.RowIter
+	var mode = unknownMode
 	if useInMemoryJoins || inMemorySession {
-		r, err := j.Right.RowIter(ctx)
-		if err != nil {
-			span.Finish()
-			return nil, err
-		}
+		mode = memoryMode
+	}
 
-		iter = &innerJoinMemoryIter{
-			l:    l,
-			r:    r,
-			ctx:  ctx,
-			cond: j.Cond,
-		}
-	} else {
-		iter = &innerJoinIter{
-			l:    l,
-			rp:   j.Right,
-			ctx:  ctx,
-			cond: j.Cond,
-		}
+	iter := &innerJoinIter{
+		l:    l,
+		rp:   j.Right,
+		ctx:  ctx,
+		cond: j.Cond,
+		mode: mode,
 	}
 
 	return sql.NewSpanIter(span, iter), nil
@@ -156,6 +183,25 @@ func (j *InnerJoin) TransformExpressions(f sql.TransformExprFunc) (sql.Node, err
 	return NewInnerJoin(j.Left, j.Right, cond), nil
 }
 
+// innerJoinMode defines the mode in which an inner join will be performed.
+type innerJoinMode byte
+
+const (
+	// unknownMode is the default mode. It will start iterating without really
+	// knowing in which mode it will end up computing the inner join. If it
+	// iterates the right side fully one time and so far it fits in memory,
+	// then it will switch to memory mode. Otherwise, if at some point during
+	// this first iteration it finds that it does not fit in memory, will
+	// switch to multipass mode.
+	unknownMode innerJoinMode = iota
+	// memoryMode computes all the inner join directly in memory iterating each
+	// side of the join exactly once.
+	memoryMode
+	// multipassMode computes the inner join by iterating the left side once,
+	// and the right side one time for each row in the left side.
+	multipassMode
+)
+
 type innerJoinIter struct {
 	l    sql.RowIter
 	rp   rowIterProvider
@@ -164,118 +210,140 @@ type innerJoinIter struct {
 	cond sql.Expression
 
 	leftRow sql.Row
-}
 
-func (i *innerJoinIter) Next() (sql.Row, error) {
-	for {
-		if i.leftRow == nil {
-			r, err := i.l.Next()
-			if err != nil {
-				return nil, err
-			}
+	// used to compute in-memory
+	mode  innerJoinMode
+	right []sql.Row
+	pos   int
+}
 
-			i.leftRow = r
+func (i *innerJoinIter) loadLeft() error {
+	if i.leftRow == nil {
+		r, err := i.l.Next()
+		if err != nil {
+			return err
 		}
 
-		if i.r == nil {
-			iter, err := i.rp.RowIter(i.ctx)
-			if err != nil {
-				return nil, err
-			}
+		i.leftRow = r
+	}
 
-			i.r = iter
-		}
+	return nil
+}
 
-		rightRow, err := i.r.Next()
-		if err == io.EOF {
-			i.r = nil
-			i.leftRow = nil
-			continue
+func (i *innerJoinIter) loadRightInMemory() error {
+	iter, err := i.rp.RowIter(i.ctx)
+	if err != nil {
+		return err
+	}
+
+	i.right, err = sql.RowIterToRows(iter)
+	if err != nil {
+		return err
+	}
+
+	if len(i.right) == 0 {
+		return io.EOF
+	}
+
+	return nil
+}
+
+func (i *innerJoinIter) fitsInMemory() bool {
+	var maxMemory uint64
+	_, v := i.ctx.Session.Get(memoryThresholdSessionVar)
+	if n, ok := v.(int64); ok {
+		maxMemory = uint64(n)
+	} else {
+		maxMemory = maxMemoryJoin
+	}
+
+	if maxMemory <= 0 {
+		return true
+	}
+
+	var ms runtime.MemStats
+	runtime.ReadMemStats(&ms)
+
+	return (ms.HeapInuse + ms.StackInuse) < maxMemory
+}
+
+func (i *innerJoinIter) loadRight() (row sql.Row, skip bool, err error) {
+	if i.mode == memoryMode {
+		if len(i.right) == 0 {
+			if err := i.loadRightInMemory(); err != nil {
+				return nil, false, err
+			}
 		}
 
-		if err != nil {
-			return nil, err
+		if i.pos >= len(i.right) {
+			i.leftRow = nil
+			i.pos = 0
+			return nil, true, nil
 		}
 
-		var row = make(sql.Row, len(i.leftRow)+len(rightRow))
-		copy(row, i.leftRow)
-		copy(row[len(i.leftRow):], rightRow)
+		row := i.right[i.pos]
+		i.pos++
+		return row, false, nil
+	}
 
-		v, err := i.cond.Eval(i.ctx, row)
+	if i.r == nil {
+		iter, err := i.rp.RowIter(i.ctx)
 		if err != nil {
-			return nil, err
+			return nil, false, err
 		}
 
-		if v == true {
-			return row, nil
-		}
+		i.r = iter
 	}
-}
 
-func (i *innerJoinIter) Close() error {
-	if err := i.l.Close(); err != nil {
-		if i.r != nil {
-			_ = i.r.Close()
+	rightRow, err := i.r.Next()
+	if err != nil {
+		if err == io.EOF {
+			i.r = nil
+			i.leftRow = nil
+
+			// If we got to this point and the mode is still unknown it means
+			// the right side fits in memory, so the mode changes to memory
+			// inner join.
+			if i.mode == unknownMode {
+				i.mode = memoryMode
+			}
+
+			return nil, true, nil
 		}
-		return err
+		return nil, false, err
 	}
 
-	if i.r != nil {
-		return i.r.Close()
+	if i.mode == unknownMode {
+		if !i.fitsInMemory() {
+			i.right = nil
+			i.mode = multipassMode
+		} else {
+			i.right = append(i.right, rightRow)
+		}
 	}
 
-	return nil
-}
-
-type innerJoinMemoryIter struct {
-	l       sql.RowIter
-	r       sql.RowIter
-	ctx     *sql.Context
-	cond    sql.Expression
-	pos     int
-	leftRow sql.Row
-	right   []sql.Row
+	return rightRow, false, err
 }
 
-func (i *innerJoinMemoryIter) Next() (sql.Row, error) {
+func (i *innerJoinIter) Next() (sql.Row, error) {
 	for {
-		if i.leftRow == nil {
-			r, err := i.l.Next()
-			if err != nil {
-				return nil, err
-			}
-
-			i.leftRow = r
+		if err := i.loadLeft(); err != nil {
+			return nil, err
 		}
 
-		if i.r != nil {
-			for {
-				row, err := i.r.Next()
-				if err != nil {
-					if err == io.EOF {
-						break
-					}
-					return nil, err
-				}
-
-				i.right = append(i.right, row)
-			}
-			i.r = nil
+		rightRow, skip, err := i.loadRight()
+		if err != nil {
+			return nil, err
 		}
 
-		if i.pos >= len(i.right) {
-			i.pos = 0
-			i.leftRow = nil
+		if skip {
 			continue
 		}
 
-		rightRow := i.right[i.pos]
 		var row = make(sql.Row, len(i.leftRow)+len(rightRow))
 		copy(row, i.leftRow)
 		copy(row[len(i.leftRow):], rightRow)
 
-		i.pos++
-
 		v, err := i.cond.Eval(i.ctx, row)
 		if err != nil {
 			return nil, err
@@ -287,7 +355,7 @@ func (i *innerJoinMemoryIter) Next() (sql.Row, error) {
 	}
 }
 
-func (i *innerJoinMemoryIter) Close() error {
+func (i *innerJoinIter) Close() error {
 	if err := i.l.Close(); err != nil {
 		if i.r != nil {
 			_ = i.r.Close()
@@ -299,5 +367,7 @@ func (i *innerJoinMemoryIter) Close() error {
 		return i.r.Close()
 	}
 
+	i.right = nil
+
 	return nil
 }