plan: compute all inner joins in memory if they fit

Fixes src-d#577

Because we do not have a way to estimate the cost of each side of
a join, it is really difficult to know when we can compute one in
memory. But not doing so, causes inner joins to be painfully slow,
as one of the branches is iterated multiple times.

This PR addresses this by ensuring that if the right branch of the
inner join fits in memory, it will be computed in memory even if
the in-memory mode has not been activated by the user.

An user can set the maximum threshold of memory the gitbase server
can use before considering the joins should not be performed in
memory using the `MAX_MEMORY_INNER_JOIN` environment variable or
the `max_memory_joins` session variable specifying the number of
bytes. The default value for this is the half of the available
physical memory on the operating system.

Because previously we had two iterators: `innerJoinIter` and
`innerJoinMemoryIter`, and now `innerJoinIter` must be able to do
the join in memory, `innerJoinMemoryIter` has been removed and
`innerJoinIter` replaced with a version that can work with three
modes:
- `unknownMode` we don't know yet how to perform the join, so keep
iterating until we can find out. By the end of the first full pass
over the right branch `unknownMode` will either switch to
`multipassMode` or `memoryMode`.
- `memoryMode` which computes the rest of the join in memory. The
iterator can have this mode before starting iterating if the user
activated the in memory join via session or environment vars, in
which case it will load all the right side on memory before doing
any further iteration. Instead, if the iterator started in
`unknownMode` and switched to this mode, it's guaranteed to already
have loaded all the right side. From that point on, they work in
exactly the same way.
- `multipassMode`, which was the previous default mode. Iterate the
right side of the join for each row in the left side. More expensive,
but less memory consuming. The iterator can not start in this mode,
and can only be switched to it from `unknownMode` in case the
memory used by the gitbase server exceeds the maximum amount of memory
either set by the user or by default.

Signed-off-by: Miguel Molina <[email protected]>

Author: erizocosmico

Diff for: 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
 }