internal/function: add cache to commit_has_blob

This commit adds an LRU cache to commit_has_blob, which avoids hitting disk many times, drastically improving performance of the UDF.

Before:

BenchmarkCommitHasBlob/commit_has_blob-4                     500           2605190 ns/op         1143268 B/op       1768 allocs/op
PASS
ok      github.com/src-d/gitquery/internal/function     2.105s

After:

BenchmarkCommitHasBlob/commit_has_blob-4                 1000000              1013 ns/op             240 B/op          7 allocs/op
PASS
ok      github.com/src-d/gitquery/internal/function     1.559s

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

Author: erizocosmico

internal/function/commit_has_blob.go

@@ -4,6 +4,8 @@ import (
 	"fmt"
 	"io"
 
+	"github.com/hashicorp/golang-lru"
+
 	"github.com/src-d/gitquery"
 	"gopkg.in/src-d/go-git.v4/plumbing"
 	"gopkg.in/src-d/go-git.v4/plumbing/object"
@@ -14,13 +16,18 @@ import (
 type CommitHasBlob struct {
 	commitHash sql.Expression
 	blob       sql.Expression
+	cache      *lru.TwoQueueCache
 }
 
+const commitHasBlobCacheSize = 200
+
 // NewCommitHasBlob creates a new commit_has_blob function.
 func NewCommitHasBlob(commitHash, blob sql.Expression) sql.Expression {
+	cache, _ := lru.New2Q(commitHasBlobCacheSize)
 	return &CommitHasBlob{
 		commitHash: commitHash,
 		blob:       blob,
+		cache:      cache,
 	}
 }
 
@@ -71,10 +78,18 @@ func (f *CommitHasBlob) Eval(ctx *sql.Context, row sql.Row) (interface{}, error)
 	)
 }
 
+type commitBlobKey struct {
+	commit, blob plumbing.Hash
+}
+
 func (f *CommitHasBlob) commitHasBlob(
 	pool *gitquery.RepositoryPool,
 	commitHash, blob plumbing.Hash,
 ) (bool, error) {
+	if val, ok := f.cache.Get(commitBlobKey{commitHash, blob}); ok {
+		return val.(bool), nil
+	}
+
 	iter, err := pool.RepoIter()
 	if err != nil {
 		return false, err
@@ -105,23 +120,29 @@ func (f *CommitHasBlob) commitHasBlob(
 			return false, err
 		}
 
-		contained, err := hashInTree(blob, tree)
+		contained, err := f.hashInTree(blob, commitHash, tree)
 		if err != nil {
 			return false, err
 		}
 
 		if contained {
 			return true, nil
 		}
+		f.cache.Add(commitBlobKey{commitHash, blob}, false)
 	}
 
 	return false, nil
 }
 
-func hashInTree(hash plumbing.Hash, tree *object.Tree) (bool, error) {
+func (f *CommitHasBlob) hashInTree(
+	hash plumbing.Hash,
+	commit plumbing.Hash,
+	tree *object.Tree,
+) (bool, error) {
 	var contained bool
-	err := tree.Files().ForEach(func(f *object.File) error {
-		if f.Blob.Hash == hash {
+	err := tree.Files().ForEach(func(fi *object.File) error {
+		f.cache.Add(commitBlobKey{commit, fi.Blob.Hash}, true)
+		if fi.Blob.Hash == hash {
 			contained = true
 			return io.EOF
 		}

internal/function/commit_has_blob_test.go

@@ -55,3 +55,44 @@ func TestCommitHasBlob(t *testing.T) {
 		})
 	}
 }
+
+func BenchmarkCommitHasBlob(b *testing.B) {
+	require.NoError(b, fixtures.Init())
+	defer func() {
+		require.NoError(b, fixtures.Clean())
+	}()
+
+	f := NewCommitHasBlob(
+		expression.NewGetField(0, sql.Text, "commit_hash", true),
+		expression.NewGetField(1, sql.Text, "blob_hash", true),
+	)
+
+	pool := gitquery.NewRepositoryPool()
+	for _, f := range fixtures.ByTag("worktree") {
+		pool.AddGit(f.Worktree().Root())
+	}
+
+	session := gitquery.NewSession(&pool)
+	ctx := sql.NewContext(context.TODO(), session)
+
+	rows := []sql.Row{
+		// blob is not on commit
+		sql.NewRow("35e85108805c84807bc66a02d91535e1e24b38b9", "9dea2395f5403188298c1dabe8bdafe562c491e3"),
+		// blob is on commit
+		sql.NewRow("6ecf0ef2c2dffb796033e5a02219af86ec6584e5", "9dea2395f5403188298c1dabe8bdafe562c491e3"),
+	}
+
+	b.Run("commit_has_blob", func(b *testing.B) {
+		require := require.New(b)
+
+		for i := 0; i < b.N; i++ {
+			val, err := f.Eval(ctx, rows[i%2])
+			require.NoError(err)
+			if i%2 == 1 {
+				require.Equal(true, val)
+			} else {
+				require.Equal(false, val)
+			}
+		}
+	})
+}

internal/function/commit_has_tree.go

@@ -20,7 +20,7 @@ import (
 // or not.
 type CommitHasTree struct {
 	expression.BinaryExpression
-	cache *lru.ARCCache
+	cache *lru.TwoQueueCache
 }
 
 // TODO: set as config
@@ -30,7 +30,7 @@ const commitHasTreeCacheSize = 100
 func NewCommitHasTree(commit, tree sql.Expression) sql.Expression {
 	// NewARC can only fail if size is negative, and we know it is not,
 	// so it is safe to ignore the error here.
-	cache, _ := lru.NewARC(commitHasTreeCacheSize)
+	cache, _ := lru.New2Q(commitHasTreeCacheSize)
 	return &CommitHasTree{expression.BinaryExpression{
 		Left:  commit,
 		Right: tree,
@@ -82,7 +82,7 @@ func (f *CommitHasTree) Eval(ctx *sql.Context, row sql.Row) (interface{}, error)
 	commitHash := plumbing.NewHash(left.(string))
 	treeHash := plumbing.NewHash(right.(string))
 
-	if val, ok := f.cache.Get(cacheKey{commitHash, treeHash}); ok {
+	if val, ok := f.cache.Get(commitTreeKey{commitHash, treeHash}); ok {
 		return val.(bool), nil
 	}
 
@@ -110,7 +110,7 @@ func (f *CommitHasTree) Eval(ctx *sql.Context, row sql.Row) (interface{}, error)
 	}
 }
 
-type cacheKey struct {
+type commitTreeKey struct {
 	commit plumbing.Hash
 	tree   plumbing.Hash
 }
@@ -124,7 +124,7 @@ func (f *CommitHasTree) commitHasTree(
 		return false, err
 	}
 
-	f.cache.Add(cacheKey{commitHash, commit.TreeHash}, true)
+	f.cache.Add(commitTreeKey{commitHash, commit.TreeHash}, true)
 
 	if commit.TreeHash == treeHash {
 		return true, nil
@@ -151,7 +151,7 @@ func (f *CommitHasTree) treeInEntries(
 
 	for {
 		if len(stack) == 0 {
-			f.cache.Add(cacheKey{commitHash, hash}, false)
+			f.cache.Add(commitTreeKey{commitHash, hash}, false)
 			return false, nil
 		}
 
@@ -164,7 +164,7 @@ func (f *CommitHasTree) treeInEntries(
 		entry := frame.entries[frame.pos]
 		frame.pos++
 		if entry.Mode == filemode.Dir {
-			f.cache.Add(cacheKey{commitHash, entry.Hash}, true)
+			f.cache.Add(commitTreeKey{commitHash, entry.Hash}, true)
 			if entry.Hash == hash {
 				return true, nil
 			}