storage: Make Put handle WriteTooOldError more like CPut

CPut (used by INSERT) and Put (used by UPDATE) previously handled
WriteTooOldErrors differently: CPut would return them immediately
while Put would defer them to the end of the transaction. This commit
makes Put work (mostly) like CPut.

The major benefit of this change is that if the first statement of a
transaction is an UPDATE, it will no longer be possible for it to
return transaction-retry errors due to SQL-level automatic retries.
(This was already true for INSERT. Note that an ON CONFLICT DO UPDATE
clause acts like a standalone UPDATE in this respect)

A potential downside is that UPDATE-heavy workloads experiencing high
contention on many keys may have worse performance (up to O(n^2)).

Updates #38591

Release note (sql change): The first statement of a transaction will
no longer return a transaction-retry error if it is an UPDATE or
DELETE (this was already true for INSERT).

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 truee 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 {