repo: big performance improvement for WriteContent with repo server (#1182)

* repo: big performance improvement for WriteContent with repo server

When re-uploading previously snapshotted directory we fetch directory
content `k<hash>` and very frequently end up writing the exact same
content. By caching last N content IDs we can avoid costly round-trip
to the server since we know that content ID was present in the session.

Also added small number of asynchronous writes, which also helps with
upload performance. Background writes are awaited before Flush().

Performance when snapshotting lots of small files (source code):

31.9 GB files:471205 dirs:75817, warm cache
Before: 260s
After: 55s (4-5x faster)

* fixed tests

internal/server/server_test.go

@@ -4,7 +4,6 @@ import (
 	"context"
 	"crypto/sha256"
 	"encoding/hex"
-	"errors"
 	"io/ioutil"
 	"net/http"
 	"net/http/httptest"
@@ -12,6 +11,7 @@ import (
 	"time"
 
 	"github.com/google/go-cmp/cmp"
+	"github.com/pkg/errors"
 	"github.com/stretchr/testify/require"
 
 	"github.com/kopia/kopia/internal/apiclient"
@@ -211,18 +211,22 @@ func remoteRepositoryTest(ctx context.Context, t *testing.T, rep repo.Repository
 
 		result = mustWriteObject(ctx, t, w, written)
 
+		require.NoError(t, w.Flush(ctx))
+
 		if uploaded == 0 {
-			t.Fatalf("did not report uploaded bytes")
+			return errors.Errorf("did not report uploaded bytes")
 		}
 
 		uploaded = 0
 		result2 := mustWriteObject(ctx, t, w, written)
+		require.NoError(t, w.Flush(ctx))
+
 		if uploaded != 0 {
-			t.Fatalf("unexpected upload when writing duplicate object")
+			return errors.Errorf("unexpected upload when writing duplicate object")
 		}
 
 		if result != result2 {
-			t.Fatalf("two identical object with different IDs: %v vs %v", result, result2)
+			return errors.Errorf("two identical object with different IDs: %v vs %v", result, result2)
 		}
 
 		// verify data is read back the same.
@@ -237,7 +241,7 @@ func remoteRepositoryTest(ctx context.Context, t *testing.T, rep repo.Repository
 
 		_, err = ow.Result()
 		if err == nil {
-			t.Fatalf("unexpected success writing object with 'm' prefix")
+			return errors.Errorf("unexpected success writing object with 'm' prefix")
 		}
 
 		manifestID, err = snapshot.SaveSnapshot(ctx, w, &snapshot.Manifest{

repo/grpc_repository_client.go

@@ -12,6 +12,7 @@ import (
 	"time"
 
 	"github.com/pkg/errors"
+	"golang.org/x/sync/errgroup"
 	"google.golang.org/grpc"
 	"google.golang.org/grpc/credentials"
 
@@ -34,6 +35,16 @@ import (
 // defined by supported splitters.
 const MaxGRPCMessageSize = 20 << 20
 
+const (
+	// number of asynchronous content writes per session.
+	grpcAsyncWritesPerSession = 4
+
+	// size of per-session cache of content IDs that were previously read
+	// helps avoid round trip to the server to write the same content since we know it already exists
+	// this greatly helps with performance of incremental snapshots.
+	numRecentReadsToCache = 1024
+)
+
 var errShouldRetry = errors.New("should retry")
 
 func errNoSessionResponse() error {
@@ -56,6 +67,9 @@ type grpcRepositoryClient struct {
 	// how many times we tried to establish inner session
 	innerSessionAttemptCount int
 
+	asyncWritesSemaphore chan struct{}
+	asyncWritesWG        errgroup.Group
+
 	h                                hashing.HashFunc
 	objectFormat                     object.Format
 	serverSupportsContentCompression bool
@@ -63,6 +77,8 @@ type grpcRepositoryClient struct {
 	omgr                             *object.Manager
 
 	contentCache *cache.PersistentCache
+
+	recent recentlyRead
 }
 
 type grpcInnerSession struct {
@@ -364,6 +380,10 @@ func (r *grpcRepositoryClient) Refresh(ctx context.Context) error {
 }
 
 func (r *grpcRepositoryClient) Flush(ctx context.Context) error {
+	if err := r.asyncWritesWG.Wait(); err != nil {
+		return errors.Wrap(err, "error waiting for async writes")
+	}
+
 	_, err := r.inSessionWithoutRetry(ctx, func(ctx context.Context, sess *grpcInnerSession) (interface{}, error) {
 		return false, sess.Flush(ctx)
 	})
@@ -502,8 +522,7 @@ func unhandledSessionResponse(resp *apipb.SessionResponse) error {
 }
 
 func (r *grpcRepositoryClient) GetContent(ctx context.Context, contentID content.ID) ([]byte, error) {
-	// nolint:wrapcheck
-	return r.contentCache.GetOrLoad(ctx, string(contentID), func() ([]byte, error) {
+	b, err := r.contentCache.GetOrLoad(ctx, string(contentID), func() ([]byte, error) {
 		v, err := r.maybeRetry(ctx, func(ctx context.Context, sess *grpcInnerSession) (interface{}, error) {
 			return sess.GetContent(ctx, contentID)
 		})
@@ -513,6 +532,13 @@ func (r *grpcRepositoryClient) GetContent(ctx context.Context, contentID content
 
 		return v.([]byte), nil
 	})
+
+	if err == nil && contentID.HasPrefix() {
+		r.recent.add(contentID)
+	}
+
+	// nolint:wrapcheck
+	return b, err
 }
 
 func (r *grpcInnerSession) GetContent(ctx context.Context, contentID content.ID) ([]byte, error) {
@@ -539,19 +565,11 @@ func (r *grpcRepositoryClient) SupportsContentCompression() bool {
 	return r.serverSupportsContentCompression
 }
 
-func (r *grpcRepositoryClient) WriteContent(ctx context.Context, data []byte, prefix content.ID, comp compression.HeaderID) (content.ID, error) {
-	if err := content.ValidatePrefix(prefix); err != nil {
-		return "", errors.Wrap(err, "invalid prefix")
-	}
-
-	var hashOutput [128]byte
-
-	contentID := prefix + content.ID(hex.EncodeToString(r.h(hashOutput[:0], data)))
-
+func (r *grpcRepositoryClient) doWrite(ctx context.Context, contentID content.ID, data []byte, prefix content.ID, comp compression.HeaderID) error {
 	// avoid uploading the content body if it already exists.
 	if _, err := r.ContentInfo(ctx, contentID); err == nil {
 		// content already exists
-		return contentID, nil
+		return nil
 	}
 
 	r.opt.OnUpload(int64(len(data)))
@@ -560,7 +578,7 @@ func (r *grpcRepositoryClient) WriteContent(ctx context.Context, data []byte, pr
 		return sess.WriteContent(ctx, data, prefix, comp)
 	})
 	if err != nil {
-		return "", err
+		return err
 	}
 
 	if prefix != "" {
@@ -568,7 +586,47 @@ func (r *grpcRepositoryClient) WriteContent(ctx context.Context, data []byte, pr
 		r.contentCache.Put(ctx, string(contentID), data)
 	}
 
-	return v.(content.ID), nil
+	if v.(content.ID) != contentID {
+		return errors.Errorf("server returned different content ID: %v vs %v", v.(content.ID), contentID)
+	}
+
+	return nil
+}
+
+func (r *grpcRepositoryClient) WriteContent(ctx context.Context, data []byte, prefix content.ID, comp compression.HeaderID) (content.ID, error) {
+	if err := content.ValidatePrefix(prefix); err != nil {
+		return "", errors.Wrap(err, "invalid prefix")
+	}
+
+	// we will be writing asynchronously and server will reject this write, fail early.
+	if prefix == manifest.ContentPrefix {
+		return "", errors.Errorf("writing manifest contents not allowed")
+	}
+
+	var hashOutput [128]byte
+
+	contentID := prefix + content.ID(hex.EncodeToString(r.h(hashOutput[:0], data)))
+
+	if r.recent.exists(contentID) {
+		return contentID, nil
+	}
+
+	// acquire semaphore
+	r.asyncWritesSemaphore <- struct{}{}
+
+	// clone so that caller can reuse the buffer
+	data = append([]byte(nil), data...)
+
+	r.asyncWritesWG.Go(func() error {
+		defer func() {
+			// release semaphore
+			<-r.asyncWritesSemaphore
+		}()
+
+		return r.doWrite(ctxutil.Detach(ctx), contentID, data, prefix, comp)
+	})
+
+	return contentID, nil
 }
 
 func (r *grpcInnerSession) WriteContent(ctx context.Context, data []byte, prefix content.ID, comp compression.HeaderID) (content.ID, error) {
@@ -758,13 +816,14 @@ func newGRPCAPIRepositoryForConnection(ctx context.Context, conn *grpc.ClientCon
 	}
 
 	rr := &grpcRepositoryClient{
-		connRefCount:       connRefCount,
-		conn:               conn,
-		cliOpts:            cliOpts,
-		transparentRetries: transparentRetries,
-		opt:                opt,
-		isReadOnly:         cliOpts.ReadOnly,
-		contentCache:       contentCache,
+		connRefCount:         connRefCount,
+		conn:                 conn,
+		cliOpts:              cliOpts,
+		transparentRetries:   transparentRetries,
+		opt:                  opt,
+		isReadOnly:           cliOpts.ReadOnly,
+		contentCache:         contentCache,
+		asyncWritesSemaphore: make(chan struct{}, grpcAsyncWritesPerSession),
 	}
 
 	v, err := rr.inSessionWithoutRetry(ctx, func(ctx context.Context, sess *grpcInnerSession) (interface{}, error) {

repo/recently_read.go

@@ -0,0 +1,46 @@
+package repo
+
+import (
+	"sync"
+
+	"github.com/kopia/kopia/repo/content"
+)
+
+type recentlyRead struct {
+	mu          sync.Mutex
+	contentList []content.ID
+	next        int
+	contentSet  map[content.ID]struct{}
+}
+
+func (r *recentlyRead) add(contentID content.ID) {
+	if r == nil {
+		return
+	}
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	if r.contentSet == nil {
+		r.contentList = make([]content.ID, numRecentReadsToCache)
+		r.contentSet = make(map[content.ID]struct{})
+	}
+
+	delete(r.contentSet, r.contentList[r.next])
+	r.contentList[r.next] = contentID
+	r.contentSet[contentID] = struct{}{}
+	r.next = (r.next + 1) % len(r.contentList)
+}
+
+func (r *recentlyRead) exists(contentID content.ID) bool {
+	if r == nil {
+		return false
+	}
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	_, ok := r.contentSet[contentID]
+
+	return ok
+}