Use all the processors to read rows from the pool

Signed-off-by: Javi Fontan <[email protected]>

Author: jfontan

blobs.go

@@ -45,7 +45,7 @@ func (r blobsTable) RowIter() (sql.RowIter, error) {
 		return nil, err
 	}
 
-	return &rowRepoIter, nil
+	return rowRepoIter, nil
 }
 
 func (blobsTable) Children() []sql.Node {
@@ -56,15 +56,15 @@ type blobIter struct {
 	iter *object.BlobIter
 }
 
-func (i *blobIter) InitRepository(repo Repository) error {
+func (i *blobIter) NewIterator(
+	repo *Repository) (RowRepoIterImplementation, error) {
+
 	iter, err := repo.Repo.BlobObjects()
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	i.iter = iter
-
-	return nil
+	return &blobIter{iter: iter}, nil
 }
 
 func (i *blobIter) Next() (sql.Row, error) {

commits.go

@@ -51,7 +51,7 @@ func (r commitsTable) RowIter() (sql.RowIter, error) {
 		return nil, err
 	}
 
-	return &rowRepoIter, nil
+	return rowRepoIter, nil
 }
 
 func (commitsTable) Children() []sql.Node {
@@ -62,15 +62,15 @@ type commitIter struct {
 	iter object.CommitIter
 }
 
-func (i *commitIter) InitRepository(repo Repository) error {
+func (i *commitIter) NewIterator(
+	repo *Repository) (RowRepoIterImplementation, error) {
+
 	iter, err := repo.Repo.CommitObjects()
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	i.iter = iter
-
-	return nil
+	return &commitIter{iter: iter}, nil
 }
 
 func (i *commitIter) Next() (sql.Row, error) {

objects.go

@@ -45,7 +45,7 @@ func (r objectsTable) RowIter() (sql.RowIter, error) {
 		return nil, err
 	}
 
-	return &rowRepoIter, nil
+	return rowRepoIter, nil
 }
 
 func (objectsTable) Children() []sql.Node {
@@ -56,15 +56,15 @@ type objectIter struct {
 	iter *object.ObjectIter
 }
 
-func (i *objectIter) InitRepository(repo Repository) error {
+func (i *objectIter) NewIterator(
+	repo *Repository) (RowRepoIterImplementation, error) {
+
 	iter, err := repo.Repo.Objects()
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	i.iter = iter
-
-	return nil
+	return &objectIter{iter: iter}, nil
 }
 
 func (i *objectIter) Next() (sql.Row, error) {

references.go

@@ -52,7 +52,7 @@ func (r referencesTable) RowIter() (sql.RowIter, error) {
 		return nil, err
 	}
 
-	return &rowRepoIter, nil
+	return rowRepoIter, nil
 }
 
 func (referencesTable) Children() []sql.Node {
@@ -63,15 +63,15 @@ type referenceIter struct {
 	iter storer.ReferenceIter
 }
 
-func (i *referenceIter) InitRepository(repo Repository) error {
+func (i *referenceIter) NewIterator(
+	repo *Repository) (RowRepoIterImplementation, error) {
+
 	iter, err := repo.Repo.References()
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	i.iter = iter
-
-	return nil
+	return &referenceIter{iter: iter}, nil
 }
 
 func (i *referenceIter) Next() (sql.Row, error) {

repositories.go

@@ -2,6 +2,8 @@ package gitquery
 
 import (
 	"io"
+	"runtime"
+	"sync"
 
 	"gopkg.in/src-d/go-git.v4"
 	"gopkg.in/src-d/go-mysql-server.v0/sql"
@@ -98,92 +100,135 @@ func (i *RepositoryIter) Close() error {
 // RowRepoIterImplementation is the interface needed by each iterator
 // implementation
 type RowRepoIterImplementation interface {
-	InitRepository(Repository) error
+	NewIterator(*Repository) (RowRepoIterImplementation, error)
 	Next() (sql.Row, error)
 	Close() error
 }
 
 // RowRepoIter is used as the base to iterate over all the repositories
-// in the pool. It needs three functions that execute specific code per
-// implemented iterator.
+// in the pool
 type RowRepoIter struct {
 	repositoryIter *RepositoryIter
-	repository     *Repository
 	implementation RowRepoIterImplementation
+
+	wg    sync.WaitGroup
+	done  chan bool
+	err   chan error
+	repos chan *Repository
+	rows  chan sql.Row
 }
 
 // NewRowRepoIter initializes a new repository iterator.
 //
 // * pool: is a RepositoryPool we want to iterate
 // * impl: implementation with RowRepoIterImplementation interface
-//     * InitRepository: called when a new repository is about to be iterated,
-//         initialize its iterator there
+//     * NewIterator: called when a new repository is about to be iterated,
+//         returns a new RowRepoIterImplementation
 //     * Next: called for each row
 //     * Close: called when a repository finished iterating
 func NewRowRepoIter(pool *RepositoryPool,
-	impl RowRepoIterImplementation) (RowRepoIter, error) {
+	impl RowRepoIterImplementation) (*RowRepoIter, error) {
 
 	rIter, err := pool.RepoIter()
 	if err != nil {
-		return RowRepoIter{}, err
+		return nil, err
 	}
 
-	repo := RowRepoIter{
+	rowRepoIter := RowRepoIter{
 		repositoryIter: rIter,
 		implementation: impl,
+		done:           make(chan bool),
+		err:            make(chan error),
+		repos:          make(chan *Repository),
+		rows:           make(chan sql.Row),
 	}
 
-	err = repo.nextRepository()
-	if err != nil {
-		return RowRepoIter{}, err
-	}
+	go rowRepoIter.fillRepoChannel()
 
-	return repo, nil
-}
+	wNum := runtime.NumCPU()
 
-// nextRepository is called to initialize the next repository. It is called
-// when the iterator is created and when the previous repository finished
-// iterating.
-func (i *RowRepoIter) nextRepository() error {
-	repo, err := i.repositoryIter.Next()
-	if err != nil {
-		return err
+	for i := 0; i < wNum; i++ {
+		rowRepoIter.wg.Add(1)
+
+		go rowRepoIter.rowReader(i)
 	}
 
-	i.repository = repo
-	err = i.implementation.InitRepository(*repo)
+	go func() {
+		rowRepoIter.wg.Wait()
+		close(rowRepoIter.rows)
+	}()
 
-	return err
+	return &rowRepoIter, nil
 }
 
-// Next gets the next row
-func (i *RowRepoIter) Next() (sql.Row, error) {
+func (i *RowRepoIter) fillRepoChannel() {
 	for {
-		row, err := i.implementation.Next()
+		repo, err := i.repositoryIter.Next()
 
 		switch err {
 		case nil:
-			return row, nil
+			i.repos <- repo
+			continue
 
 		case io.EOF:
-			i.Close()
-
-			err := i.nextRepository()
-			if err != nil {
-				return nil, err
-			}
+			close(i.repos)
+			i.err <- io.EOF
+			return
 
 		default:
-			return nil, err
+			close(i.done)
+			close(i.repos)
+			i.err <- err
+			return
+		}
+	}
+}
+
+func (i *RowRepoIter) rowReader(num int) {
+	defer i.wg.Done()
+
+	for repo := range i.repos {
+		impl, _ := i.implementation.NewIterator(repo)
+
+	loop:
+		for {
+			select {
+			case <-i.done:
+				impl.Close()
+				return
+
+			default:
+				row, err := impl.Next()
+				switch err {
+				case nil:
+					i.rows <- row
+
+				case io.EOF:
+					impl.Close()
+					break loop
+
+				default:
+					impl.Close()
+					i.err <- err
+					close(i.done)
+					return
+				}
+			}
 		}
 	}
+}
+
+// Next gets the next row
+func (i *RowRepoIter) Next() (sql.Row, error) {
+	row, ok := <-i.rows
+	if !ok {
+		return nil, <-i.err
+	}
 
-	return nil, nil
+	return row, nil
 }
 
 // Close called to close the iterator
 func (i *RowRepoIter) Close() error {
 	return i.implementation.Close()
-
-	return nil
 }

repositories_test.go

@@ -135,15 +135,15 @@ type testCommitIter struct {
 	iter object.CommitIter
 }
 
-func (d *testCommitIter) InitRepository(repo Repository) error {
-	cIter, err := repo.Repo.CommitObjects()
+func (d *testCommitIter) NewIterator(
+	repo *Repository) (RowRepoIterImplementation, error) {
+
+	iter, err := repo.Repo.CommitObjects()
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	d.iter = cIter
-
-	return nil
+	return &testCommitIter{iter: iter}, nil
 }
 
 func (d *testCommitIter) Next() (sql.Row, error) {
@@ -165,13 +165,13 @@ func TestRepositoryRowIterator(t *testing.T) {
 	path := fixtures.Basic().ByTag("worktree").One().Worktree().Root()
 
 	pool := NewRepositoryPool()
-	id, err := pool.AddGit(path)
-	require.Equal(path, id)
-	require.Nil(err)
+	max := 64
 
-	id, err = pool.AddGit(path)
-	require.Equal(path, id)
-	require.Nil(err)
+	for i := 0; i < max; i++ {
+		id, err := pool.AddGit(path)
+		require.Equal(path, id)
+		require.Nil(err)
+	}
 
 	cIter := &testCommitIter{}
 
@@ -191,5 +191,6 @@ func TestRepositoryRowIterator(t *testing.T) {
 		count++
 	}
 
-	require.Equal(9*2, count)
+	// 9 is the number of commits from the test repo
+	require.Equal(9*max, count)
 }

tags.go

@@ -50,7 +50,7 @@ func (r tagsTable) RowIter() (sql.RowIter, error) {
 		return nil, err
 	}
 
-	return &rowRepoIter, nil
+	return rowRepoIter, nil
 }
 
 func (tagsTable) Children() []sql.Node {
@@ -61,15 +61,15 @@ type tagIter struct {
 	iter *object.TagIter
 }
 
-func (i *tagIter) InitRepository(repo Repository) error {
+func (i *tagIter) NewIterator(
+	repo *Repository) (RowRepoIterImplementation, error) {
+
 	iter, err := repo.Repo.TagObjects()
 	if err != nil {
-		return err
+		return nil, err
 	}
 
-	i.iter = iter
-
-	return nil
+	return &tagIter{iter: iter}, nil
 }
 
 func (i *tagIter) Next() (sql.Row, error) {