Allow application to optionally configure storage infra clients (#719)

This PR allows applications using Tessera to (optionally) pass pre-configured instances of infrastructure clients to the drivers, enabling applications to tune those clients to the intended usage patterns and environment (e.g. https://pkg.go.dev/google.golang.org/api/option).

In this PR, we also make posix.New take a Config struct, as in the gcp and aws packages to both improve ergonomics when switching between backends, and to open the door to adding future config options to POSIX in a backwards-compatible manner.

Author: AlCutter

README_test.go

@@ -48,7 +48,7 @@ func constructStorage() {
 	ctx := context.Background()
 
 	// #region construct_example
-	driver, _ := posix.New(ctx, "/tmp/mylog")
+	driver, _ := posix.New(ctx, posix.Config{Path: "/tmp/mylog"})
 	signer := createSigner()
 
 	appender, shutdown, reader, err := tessera.NewAppender(
@@ -67,7 +67,7 @@ func constructAndUseAppender() {
 	ctx := context.Background()
 	data := []byte("hello")
 
-	driver, _ := posix.New(ctx, "/tmp/mylog")
+	driver, _ := posix.New(ctx, posix.Config{Path: "/tmp/mylog"})
 	signer := createSigner()
 
 	// #region use_appender_example

cmd/conformance/posix/main.go

@@ -69,7 +69,7 @@ func main() {
 	s, a := getSignersOrDie()
 
 	// Create the Tessera POSIX storage, using the directory from the --storage_dir flag
-	driver, err := posix.New(ctx, *storageDir)
+	driver, err := posix.New(ctx, posix.Config{Path: *storageDir})
 	if err != nil {
 		klog.Exitf("Failed to construct storage: %v", err)
 	}

cmd/examples/posix-oneshot/main.go

@@ -70,7 +70,9 @@ func main() {
 	// add all of these leaves without creating any intermediate checkpoints.
 	driver, err := posix.New(
 		ctx,
-		*storageDir,
+		posix.Config{
+			Path: *storageDir,
+		},
 	)
 	if err != nil {
 		klog.Exitf("Failed to construct storage: %v", err)

cmd/experimental/migrate/posix/main.go

@@ -63,7 +63,7 @@ func main() {
 		klog.Exitf("invalid checkpoint roothash %q: %v", bits[2], err)
 	}
 
-	driver, err := posix.New(ctx, *storageDir)
+	driver, err := posix.New(ctx, posix.Config{Path: *storageDir})
 	if err != nil {
 		klog.Exitf("Failed to create new POSIX storage driver: %v", err)
 	}

internal/witness/witness_test.go

@@ -174,7 +174,7 @@ func TestWitnessGateway_Update(t *testing.T) {
 
 func TestWitness_UpdateRequest(t *testing.T) {
 	logSignedCheckpoint, _ := loadCheckpoint(t, 9)
-	d, err := posix.New(context.Background(), "../../testdata/log/")
+	d, err := posix.New(context.Background(), posix.Config{Path: "../../testdata/log/"})
 	if err != nil {
 		t.Fatal(err)
 	}

storage/aws/aws.go

@@ -151,6 +151,9 @@ type Config struct {
 	MaxOpenConns int
 	// Maximum idle database connections in the connection pool.
 	MaxIdleConns int
+
+	// HTTPClient will be used for other HTTP requests. If unset, Tessera will use the net/http DefaultClient.
+	HTTPClient *http.Client
 }
 
 // New creates a new instance of the AWS based Storage.
@@ -221,7 +224,7 @@ func (s *Storage) newAppender(ctx context.Context, o objStore, seq sequencer, op
 		logStore:    logStore,
 		sequencer:   seq,
 		queue:       storage.NewQueue(ctx, opts.BatchMaxAge(), opts.BatchMaxSize(), seq.assignEntries),
-		newCP:       opts.CheckpointPublisher(logStore, http.DefaultClient),
+		newCP:       opts.CheckpointPublisher(logStore, s.cfg.HTTPClient),
 		treeUpdated: make(chan struct{}),
 	}
 

storage/gcp/gcp.go

@@ -123,6 +123,13 @@ type consumeFunc func(ctx context.Context, from uint64, entries []storage.Sequen
 
 // Config holds GCP project and resource configuration for a storage instance.
 type Config struct {
+	// GCSClient will be  used to interact with GCS. If unset, Tessera will create one.
+	GCSClient *gcs.Client
+	// SpannerClient will be used to interact with Spanner. If unset, Tessera will create one.
+	SpannerClient *spanner.Client
+	// HTTPClient will be used for other HTTP requests. If unset, Tessera will use the net/http DefaultClient.
+	HTTPClient *http.Client
+
 	// Bucket is the name of the GCS bucket to use for storing log state.
 	Bucket string
 	// BucketPrefix is an optional prefix to prepend to all log resource paths.
@@ -134,6 +141,9 @@ type Config struct {
 
 // New creates a new instance of the GCP based Storage.
 func New(ctx context.Context, cfg Config) (tessera.Driver, error) {
+	if cfg.HTTPClient == nil {
+		cfg.HTTPClient = http.DefaultClient
+	}
 	return &Storage{
 		cfg: cfg,
 	}, nil
@@ -192,17 +202,31 @@ func (lr *LogReader) NextIndex(ctx context.Context) (uint64, error) {
 
 // Appender creates a new tessera.Appender lifecycle object.
 func (s *Storage) Appender(ctx context.Context, opts *tessera.AppendOptions) (*tessera.Appender, tessera.LogReader, error) {
-	c, err := gcs.NewClient(ctx, gcs.WithJSONReads())
-	if err != nil {
-		return nil, nil, fmt.Errorf("failed to create GCS client: %v", err)
+	if s.cfg.GCSClient == nil {
+		var err error
+		s.cfg.GCSClient, err = gcs.NewClient(ctx, gcs.WithJSONReads())
+		if err != nil {
+			return nil, nil, fmt.Errorf("failed to create GCS client: %v", err)
+		}
 	}
 	gs := &gcsStorage{
-		gcsClient:    c,
+		gcsClient:    s.cfg.GCSClient,
 		bucket:       s.cfg.Bucket,
 		bucketPrefix: s.cfg.BucketPrefix,
 	}
 
-	seq, err := newSpannerCoordinator(ctx, s.cfg.Spanner, uint64(opts.PushbackMaxOutstanding()))
+	var err error
+	if s.cfg.SpannerClient == nil {
+		s.cfg.SpannerClient, err = spanner.NewClient(ctx, s.cfg.Spanner)
+		if err != nil {
+			return nil, nil, fmt.Errorf("failed to connect to Spanner: %v", err)
+		}
+	}
+	if err := initDB(ctx, s.cfg.Spanner); err != nil {
+		return nil, nil, fmt.Errorf("failed to verify/init Spanner schema: %v", err)
+	}
+
+	seq, err := newSpannerCoordinator(ctx, s.cfg.SpannerClient, uint64(opts.PushbackMaxOutstanding()))
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to create Spanner coordinator: %v", err)
 	}
@@ -242,7 +266,7 @@ func (s *Storage) newAppender(ctx context.Context, o objStore, seq *spannerCoord
 			return a.sequencer.nextIndex(ctx)
 		},
 	}
-	a.newCP = opts.CheckpointPublisher(reader, http.DefaultClient)
+	a.newCP = opts.CheckpointPublisher(reader, s.cfg.HTTPClient)
 
 	if err := a.init(ctx); err != nil {
 		return nil, nil, fmt.Errorf("failed to initialise log storage: %v", err)
@@ -681,18 +705,11 @@ type spannerCoordinator struct {
 
 // newSpannerCoordinator returns a new spannerSequencer struct which uses the provided
 // spanner resource name for its spanner connection.
-func newSpannerCoordinator(ctx context.Context, spannerDB string, maxOutstanding uint64) (*spannerCoordinator, error) {
-	dbPool, err := spanner.NewClient(ctx, spannerDB)
-	if err != nil {
-		return nil, fmt.Errorf("failed to connect to Spanner: %v", err)
-	}
+func newSpannerCoordinator(ctx context.Context, dbPool *spanner.Client, maxOutstanding uint64) (*spannerCoordinator, error) {
 	r := &spannerCoordinator{
 		dbPool:         dbPool,
 		maxOutstanding: maxOutstanding,
 	}
-	if err := r.initDB(ctx, spannerDB); err != nil {
-		return nil, fmt.Errorf("failed to initDB: %v", err)
-	}
 	if err := r.checkDataCompatibility(ctx); err != nil {
 		return nil, fmt.Errorf("schema is not compatible with this version of the Tessera library: %v", err)
 	}
@@ -712,7 +729,7 @@ func newSpannerCoordinator(ctx context.Context, spannerDB string, maxOutstanding
 //   - IntCoord
 //     This table coordinates integration of the batches of entries stored in
 //     Seq into the committed tree state.
-func (s *spannerCoordinator) initDB(ctx context.Context, spannerDB string) error {
+func initDB(ctx context.Context, spannerDB string) error {
 	return createAndPrepareTables(
 		ctx, spannerDB,
 		[]string{
@@ -1217,12 +1234,25 @@ func (s *gcsStorage) deleteObjectsWithPrefix(ctx context.Context, objPrefix stri
 
 // MigrationWriter creates a new GCP storage for the MigrationTarget lifecycle mode.
 func (s *Storage) MigrationWriter(ctx context.Context, opts *tessera.MigrationOptions) (migrate.MigrationWriter, tessera.LogReader, error) {
-	c, err := gcs.NewClient(ctx, gcs.WithJSONReads())
-	if err != nil {
-		return nil, nil, fmt.Errorf("failed to create GCS client: %v", err)
+	var err error
+	if s.cfg.GCSClient == nil {
+		s.cfg.GCSClient, err = gcs.NewClient(ctx, gcs.WithJSONReads())
+		if err != nil {
+			return nil, nil, fmt.Errorf("failed to create GCS client: %v", err)
+		}
+	}
+
+	if s.cfg.SpannerClient == nil {
+		s.cfg.SpannerClient, err = spanner.NewClient(ctx, s.cfg.Spanner)
+		if err != nil {
+			return nil, nil, fmt.Errorf("failed to connect to Spanner: %v", err)
+		}
+	}
+	if err := initDB(ctx, s.cfg.Spanner); err != nil {
+		return nil, nil, fmt.Errorf("failed to verify/init Spanner schema: %v", err)
 	}
 
-	seq, err := newSpannerCoordinator(ctx, s.cfg.Spanner, 0)
+	seq, err := newSpannerCoordinator(ctx, s.cfg.SpannerClient, 0)
 	if err != nil {
 		return nil, nil, fmt.Errorf("failed to create Spanner sequencer: %v", err)
 	}
@@ -1233,7 +1263,7 @@ func (s *Storage) MigrationWriter(ctx context.Context, opts *tessera.MigrationOp
 		sequencer:    seq,
 		logStore: &logResourceStore{
 			objStore: &gcsStorage{
-				gcsClient:    c,
+				gcsClient:    s.cfg.GCSClient,
 				bucket:       s.cfg.Bucket,
 				bucketPrefix: s.cfg.BucketPrefix,
 			},

storage/gcp/gcp_test.go

@@ -41,7 +41,7 @@ import (
 	"golang.org/x/mod/sumdb/note"
 )
 
-func newSpannerDB(t *testing.T) func() {
+func newSpannerDB(t *testing.T) (*spanner.Client, func()) {
 	t.Helper()
 	srv, err := spannertest.NewServer("localhost:0")
 	if err != nil {
@@ -50,16 +50,26 @@ func newSpannerDB(t *testing.T) func() {
 	if err := os.Setenv("SPANNER_EMULATOR_HOST", srv.Addr); err != nil {
 		t.Fatalf("Setenv: %v", err)
 	}
-	return srv.Close
+
+	id := "projects/p/instances/i/databases/d"
+	if err := initDB(t.Context(), id); err != nil {
+		t.Fatalf("initDB: %v", err)
+	}
+
+	c, err := spanner.NewClient(t.Context(), id)
+	if err != nil {
+		t.Fatalf("NewClient: %v", err)
+	}
+	return c, srv.Close
 
 }
 
 func TestSpannerSequencerAssignEntries(t *testing.T) {
 	ctx := context.Background()
-	close := newSpannerDB(t)
+	db, close := newSpannerDB(t)
 	defer close()
 
-	seq, err := newSpannerCoordinator(ctx, "projects/p/instances/i/databases/d", 1000)
+	seq, err := newSpannerCoordinator(ctx, db, 1000)
 	if err != nil {
 		t.Fatalf("newSpannerCoordinator: %v", err)
 	}
@@ -109,10 +119,10 @@ func TestSpannerSequencerPushback(t *testing.T) {
 		},
 	} {
 		t.Run(test.name, func(t *testing.T) {
-			close := newSpannerDB(t)
+			db, close := newSpannerDB(t)
 			defer close()
 
-			seq, err := newSpannerCoordinator(ctx, "projects/p/instances/i/databases/d", test.threshold)
+			seq, err := newSpannerCoordinator(ctx, db, test.threshold)
 			if err != nil {
 				t.Fatalf("newSpannerCoordinator: %v", err)
 			}
@@ -139,10 +149,10 @@ func TestSpannerSequencerPushback(t *testing.T) {
 
 func TestSpannerSequencerRoundTrip(t *testing.T) {
 	ctx := context.Background()
-	close := newSpannerDB(t)
+	db, close := newSpannerDB(t)
 	defer close()
 
-	s, err := newSpannerCoordinator(ctx, "projects/p/instances/i/databases/d", 1000)
+	s, err := newSpannerCoordinator(ctx, db, 1000)
 	if err != nil {
 		t.Fatalf("newSpannerCoordinator: %v", err)
 	}
@@ -191,10 +201,10 @@ func TestSpannerSequencerRoundTrip(t *testing.T) {
 
 func TestCheckDataCompatibility(t *testing.T) {
 	ctx := context.Background()
-	close := newSpannerDB(t)
+	db, close := newSpannerDB(t)
 	defer close()
 
-	s, err := newSpannerCoordinator(ctx, "projects/p/instances/i/databases/d", 1000)
+	s, err := newSpannerCoordinator(ctx, db, 1000)
 	if err != nil {
 		t.Fatalf("newSpannerCoordinator: %v", err)
 	}
@@ -381,9 +391,9 @@ func TestPublishTree(t *testing.T) {
 		},
 	} {
 		t.Run(test.name, func(t *testing.T) {
-			closeDB := newSpannerDB(t)
+			db, closeDB := newSpannerDB(t)
 			defer closeDB()
-			s, err := newSpannerCoordinator(ctx, "projects/p/instances/i/databases/d", 1000)
+			s, err := newSpannerCoordinator(ctx, db, 1000)
 			if err != nil {
 				t.Fatalf("newSpannerCoordinator: %v", err)
 			}
@@ -438,10 +448,10 @@ func TestGarbageCollect(t *testing.T) {
 	batchSize := uint64(60000)
 	integrateEvery := uint64(31234)
 
-	closeDB := newSpannerDB(t)
+	db, closeDB := newSpannerDB(t)
 	defer closeDB()
 
-	s, err := newSpannerCoordinator(ctx, "projects/p/instances/i/databases/d", batchSize)
+	s, err := newSpannerCoordinator(ctx, db, batchSize)
 	if err != nil {
 		t.Fatalf("newSpannerCoordinator: %v", err)
 	}