Add atomic writes in mem docstore

Author: sandeeppal

docstore/memdocstore/mem.go

@@ -23,6 +23,9 @@
 // Action lists are executed concurrently. Each action in an action list is executed
 // in a separate goroutine.
 //
+// memdocstore supports atomic writes. When using AtomicWrites(), all write actions
+// in the action list are executed atomically - either all succeed or all fail together.
+//
 // memdocstore calls the BeforeDo function of an ActionList once before executing the
 // actions. Its As function never returns true.
 //
@@ -198,14 +201,79 @@ func (c *collection) RunActions(ctx context.Context, actions []*driver.Action, o
 		}
 	}
 
-	beforeGets, gets, writes, _, afterGets := driver.GroupActions(actions)
+	beforeGets, gets, writes, writesTx, afterGets := driver.GroupActions(actions)
 	run(beforeGets)
 	run(gets)
 	run(writes)
+
+	// Handle atomic writes separately to ensure they are truly atomic
+	if len(writesTx) > 0 {
+		c.runAtomicWrites(ctx, writesTx, errs)
+	}
+
 	run(afterGets)
 	return driver.NewActionListError(errs)
 }
 
+// runAtomicWrites executes multiple write actions atomically.
+// All writes either succeed or all fail together.
+func (c *collection) runAtomicWrites(ctx context.Context, actions []*driver.Action, errs []error) {
+	// Stop if the context is done.
+	if ctx.Err() != nil {
+		for _, a := range actions {
+			errs[a.Index] = ctx.Err()
+		}
+		return
+	}
+
+	c.mu.Lock()
+	defer c.mu.Unlock()
+
+	// First, validate all actions and collect current documents
+	type actionInfo struct {
+		action  *driver.Action
+		current storedDoc
+		exists  bool
+	}
+
+	actionInfos := make([]actionInfo, len(actions))
+	for i, a := range actions {
+		info := &actionInfos[i]
+		info.action = a
+
+		if a.Key != nil {
+			info.current, info.exists = c.docs[a.Key]
+		}
+
+		// Check for NotFound errors
+		if !info.exists && (a.Kind == driver.Replace || a.Kind == driver.Update || a.Kind == driver.Get) {
+			for _, a2 := range actions {
+				errs[a2.Index] = gcerr.Newf(gcerr.NotFound, nil, "document with key %v does not exist", a.Key)
+			}
+			return
+		}
+
+		// Check revision conflicts
+		if err := c.checkRevision(a.Doc, info.current); err != nil {
+			for _, a2 := range actions {
+				errs[a2.Index] = err
+			}
+			return
+		}
+	}
+
+	// Now execute all actions atomically
+	for _, info := range actionInfos {
+		if err := c.executeAction(info.action, info.current, info.exists); err != nil {
+			// If any action fails, mark all actions as failed
+			for _, a2 := range actions {
+				errs[a2.Index] = err
+			}
+			return
+		}
+	}
+}
+
 // runAction executes a single action.
 func (c *collection) runAction(ctx context.Context, a *driver.Action) error {
 	// Stop if the context is done.
@@ -227,6 +295,31 @@ func (c *collection) runAction(ctx context.Context, a *driver.Action) error {
 	if !exists && (a.Kind == driver.Replace || a.Kind == driver.Update || a.Kind == driver.Get) {
 		return gcerr.Newf(gcerr.NotFound, nil, "document with key %v does not exist", a.Key)
 	}
+
+	// Check revision conflicts
+	if a.Kind != driver.Get && a.Kind != driver.Create {
+		if err := c.checkRevision(a.Doc, current); err != nil {
+			return err
+		}
+	}
+
+	// Execute the action for Get
+	if a.Kind == driver.Get {
+		// Handle Get separately since it doesn't modify the document
+		// We've already retrieved the document into current, above.
+		// Now we copy its fields into the user-provided document.
+		if err := decodeDoc(current, a.Doc, a.FieldPaths); err != nil {
+			return err
+		}
+		return nil
+	}
+
+	return c.executeAction(a, current, exists)
+}
+
+// executeAction executes a single action. Must be called with the lock held.
+// This method is shared between runAction and runAtomicWrites to eliminate code duplication.
+func (c *collection) executeAction(a *driver.Action, current storedDoc, exists bool) error {
 	switch a.Kind {
 	case driver.Create:
 		// It is an error to attempt to create an existing document.
@@ -244,9 +337,6 @@ func (c *collection) runAction(ctx context.Context, a *driver.Action) error {
 		fallthrough
 
 	case driver.Replace, driver.Put:
-		if err := c.checkRevision(a.Doc, current); err != nil {
-			return err
-		}
 		doc, err := encodeDoc(a.Doc)
 		if err != nil {
 			return err
@@ -260,15 +350,9 @@ func (c *collection) runAction(ctx context.Context, a *driver.Action) error {
 		c.docs[a.Key] = doc
 
 	case driver.Delete:
-		if err := c.checkRevision(a.Doc, current); err != nil {
-			return err
-		}
 		delete(c.docs, a.Key)
 
 	case driver.Update:
-		if err := c.checkRevision(a.Doc, current); err != nil {
-			return err
-		}
 		if err := c.update(current, a.Mods); err != nil {
 			return err
 		}
@@ -279,12 +363,6 @@ func (c *collection) runAction(ctx context.Context, a *driver.Action) error {
 			}
 		}
 
-	case driver.Get:
-		// We've already retrieved the document into current, above.
-		// Now we copy its fields into the user-provided document.
-		if err := decodeDoc(current, a.Doc, a.FieldPaths); err != nil {
-			return err
-		}
 	default:
 		return gcerr.Newf(gcerr.Internal, nil, "unknown kind %v", a.Kind)
 	}

docstore/memdocstore/mem_test.go

@@ -54,7 +54,7 @@ func (*harness) BeforeQueryTypes() []interface{} { return nil }
 
 func (*harness) RevisionsEqual(rev1, rev2 interface{}) bool { return rev1 == rev2 }
 
-func (*harness) SupportsAtomicWrites() bool { return false }
+func (*harness) SupportsAtomicWrites() bool { return true }
 
 func (*harness) Close() {}