storage: Make Put handle WriteTooOldError more like CPut

pkg/ccl/backupccl/backup.go

@@ -312,7 +312,7 @@ func ensureInterleavesIncluded(tables []*sqlbase.TableDescriptor) error {
 func allRangeDescriptors(ctx context.Context, txn *client.Txn) ([]roachpb.RangeDescriptor, error) {
 	rows, err := txn.Scan(ctx, keys.Meta2Prefix, keys.MetaMax, 0)
 	if err != nil {
-		return nil, errors.NewAssertionErrorWithWrappedErrf(err,
+		return nil, errors.Wrapf(err,
 			"unable to scan range descriptors")
 	}
 

pkg/ccl/backupccl/restore.go

@@ -352,7 +352,7 @@ func allocateTableRewrites(
 				{
 					parentDB, err := sqlbase.GetDatabaseDescFromID(ctx, txn, parentID)
 					if err != nil {
-						return errors.NewAssertionErrorWithWrappedErrf(err,
+						return errors.Wrapf(err,
 							"failed to lookup parent DB %d", errors.Safe(parentID))
 					}
 
@@ -958,7 +958,7 @@ func WriteTableDescs(
 			} else {
 				parentDB, err := sqlbase.GetDatabaseDescFromID(ctx, txn, tables[i].ParentID)
 				if err != nil {
-					return errors.NewAssertionErrorWithWrappedErrf(err,
+					return errors.Wrapf(err,
 						"failed to lookup parent DB %d", errors.Safe(tables[i].ParentID))
 				}
 				// TODO(mberhault): CheckPrivilege wants a planner.
@@ -984,7 +984,7 @@ func WriteTableDescs(
 
 		for _, table := range tables {
 			if err := table.Validate(ctx, txn, settings); err != nil {
-				return errors.NewAssertionErrorWithWrappedErrf(err,
+				return errors.Wrapf(err,
 					"validate table %d", errors.Safe(table.ID))
 			}
 		}

pkg/kv/dist_sender_server_test.go

@@ -2192,6 +2192,19 @@ func TestTxnCoordSenderRetries(t *testing.T) {
 			retryable: func(ctx context.Context, txn *client.Txn) error {
 				return txn.Put(ctx, "a", "put")
 			},
+			txnCoordRetry: true,
+		},
+		{
+			name: "deferred write too old with put",
+			afterTxnStart: func(ctx context.Context, db *client.DB) error {
+				return db.Put(ctx, "a", "put")
+			},
+			retryable: func(ctx context.Context, txn *client.Txn) error {
+				b := txn.NewBatch()
+				b.Header.DeferWriteTooOldError = true
+				b.Put("a", "put")
+				return txn.Run(ctx, b)
+			},
 			// This trivially succeeds as there are no refresh spans.
 		},
 		{
@@ -2510,6 +2523,20 @@ func TestTxnCoordSenderRetries(t *testing.T) {
 				b := txn.NewBatch()
 				b.Put("a", "put")
 				b.Put("c", "put")
+				return txn.CommitInBatch(ctx, b) // put to c will return WriteTooOldError
+			},
+			clientRetry: true,
+		},
+		{
+			name: "multi-range batch with deferred write too old",
+			afterTxnStart: func(ctx context.Context, db *client.DB) error {
+				return db.Put(ctx, "c", "value")
+			},
+			retryable: func(ctx context.Context, txn *client.Txn) error {
+				b := txn.NewBatch()
+				b.Header.DeferWriteTooOldError = true
+				b.Put("a", "put")
+				b.Put("c", "put")
 				return txn.CommitInBatch(ctx, b) // both puts will succeed, et will retry
 			},
 			// Parallel commits do not support the canForwardSerializableTimestamp
@@ -2549,6 +2576,40 @@ func TestTxnCoordSenderRetries(t *testing.T) {
 			},
 			clientRetry: true, // successful cput will still retry because of mixed success
 		},
+		{
+			name: "multi-range batch with deferred write too old and failed cput",
+			beforeTxnStart: func(ctx context.Context, db *client.DB) error {
+				return db.Put(ctx, "a", "orig")
+			},
+			afterTxnStart: func(ctx context.Context, db *client.DB) error {
+				return db.Put(ctx, "a", "value")
+			},
+			retryable: func(ctx context.Context, txn *client.Txn) error {
+				b := txn.NewBatch()
+				b.Header.DeferWriteTooOldError = true
+				b.CPut("a", "cput", "orig")
+				b.Put("c", "put")
+				return txn.CommitInBatch(ctx, b)
+			},
+			clientRetry: true, // cput with write too old requires restart
+		},
+		{
+			name: "multi-range batch with deferred write too old and successful cput",
+			beforeTxnStart: func(ctx context.Context, db *client.DB) error {
+				return db.Put(ctx, "a", "orig")
+			},
+			afterTxnStart: func(ctx context.Context, db *client.DB) error {
+				return db.Put(ctx, "a", "orig")
+			},
+			retryable: func(ctx context.Context, txn *client.Txn) error {
+				b := txn.NewBatch()
+				b.Header.DeferWriteTooOldError = true
+				b.CPut("a", "cput", "orig")
+				b.Put("c", "put")
+				return txn.CommitInBatch(ctx, b)
+			},
+			clientRetry: true, // successful cput will still retry because of mixed success
+		},
 		{
 			name: "cput within uncertainty interval",
 			beforeTxnStart: func(ctx context.Context, db *client.DB) error {

pkg/kv/txn_test.go

@@ -156,14 +156,17 @@ func TestLostUpdate(t *testing.T) {
 			t.Fatal("should experience just one restart")
 		}
 
+		var newVal string
 		if gr.Exists() && bytes.Equal(gr.ValueBytes(), []byte("hi")) {
-			if err := txn.Put(ctx, key, "correct"); err != nil {
-				t.Fatal(err)
-			}
+			newVal = "correct"
 		} else {
-			if err := txn.Put(ctx, key, "oops!"); err != nil {
-				t.Fatal(err)
-			}
+			newVal = "oops!"
+		}
+		b := txn.NewBatch()
+		b.Header.DeferWriteTooOldError = true
+		b.Put(key, newVal)
+		if err := txn.Run(ctx, b); err != nil {
+			t.Fatal(err)
 		}
 		// Verify that the WriteTooOld boolean is set on the txn.
 		proto := txn.Serialize()

pkg/roachpb/api.go

@@ -108,6 +108,13 @@ func IsReadOnly(args Request) bool {
 	return (flags&isRead) != 0 && (flags&isWrite) == 0
 }
 
+// IsReadAndWrite returns true if the request both reads and writes
+// (such as conditional puts).
+func IsReadAndWrite(args Request) bool {
+	flags := args.flags()
+	return (flags&isRead) != 0 && (flags&isWrite) != 0
+}
+
 // IsTransactional returns true if the request may be part of a
 // transaction.
 func IsTransactional(args Request) bool {