support for atomic writes for gcp (#3523)

* support for atomic writes for gcp

* docstore add mongo  and firestore as well

* mongodb to single replica replicaset

* fix test name

* fix test

Author: sandeepvinayak

docstore/awsdynamodb/dynamo_test.go

@@ -60,6 +60,10 @@ type harness struct {
 	closer func()
 }
 
+func (h *harness) SupportsAtomicWrites() bool {
+	return true
+}
+
 func newHarness(ctx context.Context, t *testing.T) (drivertest.Harness, error) {
 	t.Helper()
 

docstore/drivertest/drivertest.go

@@ -23,6 +23,7 @@ import (
 	"io"
 	"math"
 	"reflect"
+	"strings"
 	"testing"
 	"time"
 
@@ -80,6 +81,9 @@ type Harness interface {
 	// RevisionsEqual reports whether two revisions are equal.
 	RevisionsEqual(rev1, rev2 interface{}) bool
 
+	// SupportsAtomicWrites should report if a collection supports atomic writes
+	SupportsAtomicWrites() bool
+
 	// Close closes resources used by the harness.
 	Close()
 }
@@ -170,6 +174,8 @@ func RunConformanceTests(t *testing.T, newHarness HarnessMaker, ct CodecTester,
 	t.Run("BlindCodec", func(t *testing.T) { testBlindDecode(t, ct) })
 
 	t.Run("Create", func(t *testing.T) { withRevCollections(t, newHarness, testCreate) })
+	t.Run("AtomicWrites", func(t *testing.T) { withRevCollections(t, newHarness, testAtomicWrites) })
+	t.Run("AtomicWritesFail", func(t *testing.T) { withRevCollections(t, newHarness, testAtomicWritesFail) })
 	t.Run("Put", func(t *testing.T) { withRevCollections(t, newHarness, testPut) })
 	t.Run("Replace", func(t *testing.T) { withRevCollections(t, newHarness, testReplace) })
 	t.Run("Get", func(t *testing.T) { withRevCollections(t, newHarness, testGet) })
@@ -234,6 +240,9 @@ func withRevCollections(t *testing.T, newHarness HarnessMaker, f func(*testing.T
 		t.Fatal(err)
 	}
 	defer h.Close()
+	if strings.Contains(t.Name(), "AtomicWrites") && !h.SupportsAtomicWrites() {
+		t.Skip()
+	}
 
 	t.Run("StdRev", func(t *testing.T) {
 		withColl(t, h, SingleKey, func(t *testing.T, _ Harness, coll *docstore.Collection) {

docstore/gcpfirestore/fs.go

@@ -265,8 +265,9 @@ func (c *collection) RevisionField() string {
 // RunActions implements driver.RunActions.
 func (c *collection) RunActions(ctx context.Context, actions []*driver.Action, opts *driver.RunActionsOptions) driver.ActionListError {
 	errs := make([]error, len(actions))
-	beforeGets, gets, writes, _, afterGets := driver.GroupActions(actions)
+	beforeGets, gets, writes, writesTx, afterGets := driver.GroupActions(actions)
 	calls := c.buildCommitCalls(writes, errs)
+	atomicWritesCall := c.buildAtomicWritesCommitCall(writesTx, errs)
 	// runGets does not issue concurrent RPCs, so it doesn't need a throttle.
 	c.runGets(ctx, beforeGets, errs, opts)
 	t := driver.NewThrottle(c.opts.MaxOutstandingActionRPCs)
@@ -278,6 +279,15 @@ func (c *collection) RunActions(ctx context.Context, actions []*driver.Action, o
 			c.doCommitCall(ctx, call, errs, opts)
 		}()
 	}
+	// commit the atomic writes
+	if atomicWritesCall != nil && len(atomicWritesCall.actions) > 0 {
+		t.Acquire()
+		go func() {
+			defer t.Release()
+			c.doCommitCall(ctx, atomicWritesCall, errs, opts)
+		}()
+	}
+
 	t.Acquire()
 	c.runGets(ctx, gets, errs, opts)
 	t.Release()
@@ -407,6 +417,23 @@ func (c *collection) buildCommitCalls(actions []*driver.Action, errs []error) []
 	return append(pCalls, nCall)
 }
 
+// Construct a commit call with all the atomic writes
+func (c *collection) buildAtomicWritesCommitCall(actions []*driver.Action, errs []error) *commitCall {
+	atomicWritesCommitCall := &commitCall{}
+	for _, a := range actions {
+		ws, nn, err := c.actionToWrites(a)
+		if err != nil {
+			errs[a.Index] = err
+			return nil
+		} else {
+			atomicWritesCommitCall.writes = append(atomicWritesCommitCall.writes, ws...)
+			atomicWritesCommitCall.actions = append(atomicWritesCommitCall.actions, a)
+			atomicWritesCommitCall.newNames = append(atomicWritesCommitCall.newNames, nn)
+		}
+	}
+	return atomicWritesCommitCall
+}
+
 // Convert an action to one or more Firestore Write protos.
 func (c *collection) actionToWrites(a *driver.Action) ([]*pb.Write, string, error) {
 	var (

docstore/gcpfirestore/fs_test.go

@@ -33,7 +33,7 @@ import (
 
 const (
 	// projectID is the project ID that was used during the last test run using --record.
-	projectID       = "go-cloud-test-216917"
+	projectID       = "go-cloud-test-216918"
 	collectionName1 = "docstore-test-1"
 	collectionName2 = "docstore-test-2"
 	collectionName3 = "docstore-test-3"
@@ -45,6 +45,10 @@ type harness struct {
 	done   func()
 }
 
+func (h *harness) SupportsAtomicWrites() bool {
+	return true
+}
+
 func newHarness(ctx context.Context, t *testing.T) (drivertest.Harness, error) {
 	t.Helper()
 

docstore/gcpfirestore/testdata/TestConformance/ActionsOnStructNoRev.replay

@@ -52,6 +52,8 @@ func (*harness) BeforeQueryTypes() []interface{} { return nil }
 
 func (*harness) RevisionsEqual(rev1, rev2 interface{}) bool { return rev1 == rev2 }
 
+func (*harness) SupportsAtomicWrites() bool { return false }
+
 func (*harness) Close() {}
 
 func TestConformance(t *testing.T) {

docstore/gcpfirestore/testdata/TestConformance/ActionsWithCompositeID.replay

@@ -52,7 +52,7 @@ func TestConformanceCosmos(t *testing.T) {
 	}()
 
 	newHarness := func(context.Context, *testing.T) (drivertest.Harness, error) {
-		return &harness{client.Database(dbName)}, nil
+		return &harness{client.Database(dbName), true}, nil
 	}
 	drivertest.RunConformanceTests(t, newHarness, codecTester{}, []drivertest.AsTest{verifyAs{}})
 }

docstore/gcpfirestore/testdata/TestConformance/As/verify_As.replay

@@ -0,0 +1,30 @@
+version: '3.8'
+services:
+  mongo:
+    image: mongo:4.4
+    container_name: mongo
+    restart: always
+    volumes:
+      - mongo_data:/data/db
+    ports:
+      - "27017:27017"
+    command: ["--replSet", "dbrs", "--bind_ip_all"]
+
+  mongosetup:
+    image: mongo:4.4
+    container_name: mongosetup
+    depends_on:
+      - mongo
+    entrypoint: >
+      bash -c "sleep 5 &&
+      mongo --host mongo:27017 --eval '
+        rs.initiate({
+          _id: \"dbrs\",
+          members: [
+            { _id: 0, host: \"mongo:27017\"}
+          ]
+        })
+      '"
+
+volumes:
+  mongo_data:

docstore/gcpfirestore/testdata/TestConformance/As/verify_As_returns_false_when_passed_nil.replay

@@ -20,14 +20,15 @@
 set -euo pipefail
 
 echo "Starting MongoDB v4 listening on 27017..."
-docker rm -f mongo4 &> /dev/null || :
-docker run -d --name mongo4  -p 27017:27017 mongo:4 &> /dev/null
-echo "...done. Run \"docker rm -f mongo4\" to clean up the container."
+docker rm -f mongo4 mongosetup &> /dev/null || :
+docker compose -f ./docstore/mongodocstore/docker-compose.yml up --wait &> /dev/null
+sleep 3
+echo "...done. Run \"docker rm -f mongo4 mongosetup\" to clean up the container."
 echo
 
-echo "Starting MongoDB v3 listening on 27018..."
+echo "Starting MongoDB v3 listening on 27020..."
 docker rm -f mongo3 &> /dev/null || :
-docker run -d --name mongo3  -p 27018:27017 mongo:3 &> /dev/null
+docker run -d --name mongo3  -p 27020:27017 mongo:3 &> /dev/null
 echo "...done. Run \"docker rm -f mongo3\" to clean up the container."
 echo
 

docstore/gcpfirestore/testdata/TestConformance/AtomicWrites.replay

@@ -204,12 +204,15 @@ const mongoIDField = "_id"
 
 func (c *collection) RunActions(ctx context.Context, actions []*driver.Action, opts *driver.RunActionsOptions) driver.ActionListError {
 	errs := make([]error, len(actions))
-	beforeGets, gets, writes, _, afterGets := driver.GroupActions(actions)
+	beforeGets, gets, writes, writesTx, afterGets := driver.GroupActions(actions)
 	c.runGets(ctx, beforeGets, errs, opts)
 	ch := make(chan []error)
 	go func() { ch <- c.bulkWrite(ctx, writes, errs, opts) }()
+	ch2 := make(chan []error)
+	go func() { ch2 <- c.txWrite(ctx, writesTx, errs, opts) }()
 	c.runGets(ctx, gets, errs, opts)
 	writeErrs := <-ch
+	<-ch2
 	c.runGets(ctx, afterGets, errs, opts)
 	alerr := driver.NewActionListError(errs)
 	for _, werr := range writeErrs {
@@ -560,8 +563,136 @@ func (c *collection) bulkWrite(ctx context.Context, actions []*driver.Action, er
 	return reterrs
 }
 
-func (c *collection) determineDeleteErrors(ctx context.Context, models []mongo.WriteModel, actions []*driver.Action, errs []error) {
+func (c *collection) txWrite(ctx context.Context, actions []*driver.Action, errs []error, dopts *driver.RunActionsOptions) []error {
+	var (
+		models          []mongo.WriteModel
+		modelActions    []*driver.Action // corresponding action for each model
+		newIDs          []interface{}    // new IDs for Create actions, corresponding to models slice
+		revs            []string         // new revisions, corresponding to models slice
+		nDeletes        int64
+		nNonCreateWrite int64 // total operations expected from Put, Replace and Update
+	)
+
+	// all actions will fail atomically even if a single action fails
+	setErr := func(err error) {
+		for _, a := range actions {
+			errs[a.Index] = err
+		}
+	}
+
+	for _, a := range actions {
+		var m mongo.WriteModel
+		var err error
+		var newID interface{}
+		var rev string
+		switch a.Kind {
+		case driver.Create:
+			m, newID, rev, err = c.newCreateModel(a)
+		case driver.Delete:
+			m, err = c.newDeleteModel(a)
+			if err == nil {
+				nDeletes++
+			}
+		case driver.Replace, driver.Put:
+			m, rev, err = c.newReplaceModel(a, a.Kind == driver.Put)
+			if err == nil {
+				nNonCreateWrite++
+			}
+		case driver.Update:
+			m, rev, err = c.newUpdateModel(a)
+			if err == nil && m != nil {
+				nNonCreateWrite++
+			}
+		default:
+			err = gcerr.Newf(gcerr.Internal, nil, "bad action %+v", a)
+		}
+		if err != nil {
+			setErr(err)
+			return nil
+		} else if m != nil { // m can be nil for a no-op update
+			models = append(models, m)
+			modelActions = append(modelActions, a)
+			newIDs = append(newIDs, newID)
+			revs = append(revs, rev)
+		}
+	}
+	if len(models) == 0 {
+		return nil
+	}
+
+	bopts := options.BulkWrite().SetOrdered(true)
+	if dopts.BeforeDo != nil {
+		asFunc := func(target interface{}) bool {
+			switch t := target.(type) {
+			case *[]mongo.WriteModel:
+				*t = models
+			case **options.BulkWriteOptions:
+				*t = bopts
+			default:
+				return false
+			}
+			return true
+		}
+		if err := dopts.BeforeDo(asFunc); err != nil {
+			return []error{err}
+		}
+	}
+
+	client := c.coll.Database().Client()
+	session, err := client.StartSession()
+	if err != nil {
+		setErr(err)
+		return nil
+	}
+	defer session.EndSession(ctx)
+
+	callback := func(sessionCtx mongo.SessionContext) error {
+		if err := session.StartTransaction(); err != nil {
+			return err
+		}
+
+		res, err := c.coll.BulkWrite(sessionCtx, models, bopts)
+		if res.DeletedCount != nDeletes {
+			// Some Delete actions failed. It's not an error if a Delete failed because
+			// the document didn't exist, but it is an error if it failed because of a
+			// precondition mismatch. Find all the documents with revisions we tried to delete; if
+			// any are still present, that's an error.
+			if c.determineDeleteErrors(ctx, models, modelActions, errs) {
+				setErr(gcerr.Newf(gcerr.FailedPrecondition, nil,
+					"wrong revision for document to be deleted"))
+			}
+		}
+		if res.MatchedCount+res.UpsertedCount != nNonCreateWrite {
+			err = gcerr.Newf(gcerr.NotFound, nil, "some writes failed (replaced %d, upserted %d, out of total %d)", res.MatchedCount, res.UpsertedCount, nNonCreateWrite)
+		}
+		if err != nil {
+			abortTxErr := session.AbortTransaction(context.Background())
+			if abortTxErr != nil {
+				return err
+			}
+			return err
+		}
+
+		if err = session.CommitTransaction(sessionCtx); err != nil {
+			return err
+		}
+
+		return nil
+	}
+
+	err = mongo.WithSession(ctx, session, callback)
+
+	if err != nil {
+		setErr(err)
+		return nil
+	}
+	return nil
+}
+
+// determineDeleteErrors find the errors for the delete and return true if found any
+func (c *collection) determineDeleteErrors(ctx context.Context, models []mongo.WriteModel, actions []*driver.Action, errs []error) bool {
 	// TODO(jba): do this concurrently.
+	foundErr := false
 	for i, m := range models {
 		if dm, ok := m.(*mongo.DeleteOneModel); ok {
 			filter := dm.Filter.(bson.D)
@@ -580,10 +711,12 @@ func (c *collection) determineDeleteErrors(ctx context.Context, models []mongo.W
 					// revision.
 					errs[actions[i].Index] = gcerr.Newf(gcerr.FailedPrecondition, nil,
 						"wrong revision for document with ID %v", actions[i].Key)
+					foundErr = true
 				}
 			}
 		}
 	}
+	return foundErr
 }
 
 func (c *collection) newCreateModel(a *driver.Action) (*mongo.InsertOneModel, interface{}, string, error) {