feat(spanner): long running transaction clean up - disabled

spanner/README.md

@@ -31,4 +31,87 @@ row, err := client.Single().ReadRow(ctx, "Users",
 if err != nil {
 	log.Fatal(err)
 }
+```
+
+### Session Leak
+A `Client` object of the Client Library has a limit on the number of maximum sessions. For example the
+default value of `MaxOpened`, which is the maximum number of sessions allowed by the session pool in the 
+Golang Client Library is 400. You can configure these values at the time of
+creating a `Client` by passing custom `SessionPoolConfig` as part of `ClientConfig`. When all the sessions are checked
+out of the session pool, every new transaction has to wait until a session is returned to the pool.
+If a session is never returned to the pool (hence causing a session leak), the transactions will have to wait
+indefinitely and your application will be blocked.
+
+#### Common Root Causes
+The most common reason for session leaks in the Golang client library are:
+1. Not stopping a `RowIterator` that is returned by `Query`, `Read` and other methods. Always use `RowIterator.Stop()` to ensure that the `RowIterator` is always closed.
+2. Not closing a `ReadOnlyTransaction` when you no longer need it. Always call `ReadOnlyTransaction.Close()` after use, to ensure that the `ReadOnlyTransaction` is always closed.
+3. Not closing a `BatchReadOnlyTransaction` when you no longer need it. Always call `BatchReadOnlyTransaction.Close()` after use, to ensure that the `BatchReadOnlyTransaction` is always closed.
+
+As shown in the example below, the `txn.Close()` statement releases the session after it is complete.
+If you don't do `txn.Close()`, the session is not released back to the pool. The recommended way is to use `defer` as shown below,
+```go
+client, err := spanner.NewClient(ctx, "projects/P/instances/I/databases/D")
+if err != nil {
+  log.Fatal(err)
+}
+txn := client.ReadOnlyTransaction()
+defer txn.Close()
+```
+
+#### Debugging and Resolving Session Leaks
+
+##### Logging inactive transactions
+Enabled by default, the logging option shares warn logs when you have exhausted >95% of your session pool.
+This could mean two things, either you need to increase the max sessions in your session pool (as the number
+of queries run using the client side database object is greater than your session pool can serve) or you may
+have a session leak. To help debug which transactions may be causing this session leak, the logs will also contain stack traces of
+transactions which have been running longer than expected if `TrackSessionHandles` under `SessionPoolConfig` is enabled.
+
+```go
+sessionPoolConfig := spanner.SessionPoolConfig{
+    TrackSessionHandles: true,
+    InactiveTransactionRemovalOptions: spanner.InactiveTransactionRemovalOptions{
+      ActionOnInactiveTransaction: spanner.Warn,
+    },
+}
+client, err := spanner.NewClientWithConfig(
+	ctx, database, spanner.ClientConfig{SessionPoolConfig: sessionPoolConfig},
+)
+if err != nil {
+	log.Fatal(err)
+}
+defer client.Close()
+
+// Example Log message to warn presence of long running transactions
+// session <session-info> checked out of pool at <session-checkout-time> is long running due to possible session leak for goroutine
+// <Stack Trace of transaction>
+
+```
+
+##### Automatically clean inactive transactions
+When the option to automatically clean inactive transactions is enabled, the client library will automatically spot
+problematic transactions that are running for extremely long periods of time (thus causing session leaks) and close them.
+The session will be removed from the pool and be replaced by a new session. To dig deeper into which transactions are being
+closed, you can check the logs to see the stack trace of the transactions which might be causing these leaks and further
+debug them.
+
+```go
+sessionPoolConfig := spanner.SessionPoolConfig{
+    TrackSessionHandles: true,
+    InactiveTransactionRemovalOptions: spanner.InactiveTransactionRemovalOptions{
+      ActionOnInactiveTransaction: spanner.WarnAndClose,
+    },
+}
+client, err := spanner.NewClientWithConfig(
+	ctx, database, spanner.ClientConfig{SessionPoolConfig: sessionPoolConfig},
+)
+if err != nil {
+log.Fatal(err)
+}
+defer client.Close()
+
+// Example Log message for when transaction is recycled
+// session <session-info> checked out of pool at <session-checkout-time> is long running and will be removed due to possible session leak for goroutine 
+// <Stack Trace of transaction>
 ```

spanner/client.go

@@ -439,10 +439,11 @@ func (c *Client) BatchReadOnlyTransaction(ctx context.Context, tb TimestampBound
 
 	t := &BatchReadOnlyTransaction{
 		ReadOnlyTransaction: ReadOnlyTransaction{
-			tx:              tx,
-			txReadyOrClosed: make(chan struct{}),
-			state:           txActive,
-			rts:             rts,
+			tx:                       tx,
+			txReadyOrClosed:          make(chan struct{}),
+			state:                    txActive,
+			rts:                      rts,
+			isLongRunningTransaction: true,
 		},
 		ID: BatchReadOnlyTransactionID{
 			tid: tx,
@@ -473,10 +474,11 @@ func (c *Client) BatchReadOnlyTransactionFromID(tid BatchReadOnlyTransactionID)
 
 	t := &BatchReadOnlyTransaction{
 		ReadOnlyTransaction: ReadOnlyTransaction{
-			tx:              tid.tid,
-			txReadyOrClosed: make(chan struct{}),
-			state:           txActive,
-			rts:             tid.rts,
+			tx:                       tid.tid,
+			txReadyOrClosed:          make(chan struct{}),
+			state:                    txActive,
+			rts:                      tid.rts,
+			isLongRunningTransaction: true,
 		},
 		ID: tid,
 	}
@@ -558,6 +560,14 @@ func (c *Client) rwTransaction(ctx context.Context, f func(context.Context, *Rea
 		if sh == nil || sh.getID() == "" || sh.getClient() == nil {
 			// Session handle hasn't been allocated or has been destroyed.
 			sh, err = c.idleSessions.take(ctx)
+			if t != nil {
+				// Some operations (for ex BatchUpdate) can be long-running. For such operations set the isLongRunningTransaction flag to be true
+				sh.mu.Lock()
+				t.mu.Lock()
+				sh.eligibleForLongRunning = t.isLongRunningTransaction
+				t.mu.Unlock()
+				sh.mu.Unlock()
+			}
 			if err != nil {
 				// If session retrieval fails, just fail the transaction.
 				return err

spanner/client_test.go

@@ -208,6 +208,73 @@ func TestClient_Single_Read_SessionNotFound(t *testing.T) {
 	}
 }
 
+func TestClient_Single_WhenInactiveTransactionsAndSessionIsNotFoundOnBackend_RemoveSessionFromPool(t *testing.T) {
+	t.Parallel()
+	server, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{
+		SessionPoolConfig: SessionPoolConfig{
+			MinOpened: 1,
+			MaxOpened: 1,
+			InactiveTransactionRemovalOptions: InactiveTransactionRemovalOptions{
+				ActionOnInactiveTransaction: WarnAndClose,
+			},
+		},
+	})
+	defer teardown()
+	server.TestSpanner.PutExecutionTime(
+		MethodExecuteStreamingSql,
+		SimulatedExecutionTime{Errors: []error{newSessionNotFoundError("projects/p/instances/i/databases/d/sessions/s")}},
+	)
+	ctx := context.Background()
+	single := client.Single()
+	iter := single.Query(ctx, NewStatement(SelectSingerIDAlbumIDAlbumTitleFromAlbums))
+	p := client.idleSessions
+	sh := single.sh
+	// simulate session to be checked out for more than 60mins
+	sh.mu.Lock()
+	sh.checkoutTime = time.Now().Add(-time.Hour)
+	sh.mu.Unlock()
+
+	// force run task to clean up unexpected long-running sessions
+	p.removeLongRunningSessions()
+	rowCount := int64(0)
+	for {
+		// Backend throws SessionNotFoundError. Session gets replaced with new session
+		_, err := iter.Next()
+		if err == iterator.Done {
+			break
+		}
+		if err != nil {
+			t.Fatal(err)
+		}
+		rowCount++
+	}
+	// New session returns back to pool
+	iter.Stop()
+
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	if g, w := p.idleList.Len(), 1; g != w {
+		t.Fatalf("Idle Sessions in pool, count mismatch\nGot: %d\nWant: %d\n", g, w)
+	}
+	if g, w := p.numInUse, uint64(0); g != w {
+		t.Fatalf("Number of sessions currently in use mismatch\nGot: %d\nWant: %d\n", g, w)
+	}
+	if g, w := p.numOpened, uint64(1); g != w {
+		t.Fatalf("Session pool size mismatch\nGot: %d\nWant: %d\n", g, w)
+	}
+
+	sh.mu.Lock()
+	defer sh.mu.Unlock()
+	if g, w := sh.eligibleForLongRunning, false; g != w {
+		t.Fatalf("isLongRunningTransaction mismatch\nGot: %v\nWant: %v\n", g, w)
+	}
+	p.InactiveTransactionRemovalOptions.mu.Lock()
+	defer p.InactiveTransactionRemovalOptions.mu.Unlock()
+	if g, w := p.numOfLeakedSessionsRemoved, uint64(1); g != w {
+		t.Fatalf("Number of leaked sessions removed mismatch\nGot: %d\nWant: %d\n", g, w)
+	}
+}
+
 func TestClient_Single_ReadRow_SessionNotFound(t *testing.T) {
 	t.Parallel()
 
@@ -1368,6 +1435,59 @@ func TestClient_ReadWriteTransaction_SessionNotFoundOnExecuteUpdate(t *testing.T
 	}
 }
 
+func TestClient_ReadWriteTransaction_WhenLongRunningSessionCleaned_TransactionShouldFail(t *testing.T) {
+	t.Parallel()
+	_, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{
+		SessionPoolConfig: SessionPoolConfig{
+			MinOpened: 1,
+			MaxOpened: 1,
+			InactiveTransactionRemovalOptions: InactiveTransactionRemovalOptions{
+				ActionOnInactiveTransaction: WarnAndClose,
+			},
+		},
+	})
+	defer teardown()
+	ctx := context.Background()
+	p := client.idleSessions
+	msg := "session is already recycled / destroyed"
+	_, err := client.ReadWriteTransaction(ctx, func(ctx context.Context, tx *ReadWriteTransaction) error {
+		rowCount, err := tx.Update(ctx, NewStatement(UpdateBarSetFoo))
+		if err != nil {
+			return err
+		}
+		if g, w := rowCount, int64(UpdateBarSetFooRowCount); g != w {
+			return status.Errorf(codes.FailedPrecondition, "Row count mismatch\nGot: %v\nWant: %v", g, w)
+		}
+
+		// Simulate the session to be checked out for more than 60 mins.
+		// The background task cleans up this long-running session.
+		tx.sh.mu.Lock()
+		tx.sh.checkoutTime = time.Now().Add(-time.Hour)
+		if g, w := tx.sh.eligibleForLongRunning, false; g != w {
+			tx.sh.mu.Unlock()
+			return status.Errorf(codes.FailedPrecondition, "isLongRunningTransaction value mismatch\nGot: %v\nWant: %v", g, w)
+		}
+		tx.sh.mu.Unlock()
+
+		// force run task to clean up unexpected long-running sessions
+		p.removeLongRunningSessions()
+
+		// The session associated with this transaction tx has been destroyed. So the below call should fail.
+		// Eventually this means the entire transaction should not succeed.
+		_, err = tx.Update(ctx, NewStatement("UPDATE FOO SET BAR='value' WHERE ID=1"))
+		if err != nil {
+			return err
+		}
+		return nil
+	})
+	if err == nil {
+		t.Fatalf("Missing expected exception")
+	}
+	if status.Code(err) != codes.FailedPrecondition || !strings.Contains(err.Error(), msg) {
+		t.Fatalf("error mismatch\nGot: %v\nWant: %v", err, msg)
+	}
+}
+
 func TestClient_ReadWriteTransaction_SessionNotFoundOnExecuteBatchUpdate(t *testing.T) {
 	t.Parallel()
 
@@ -1401,6 +1521,65 @@ func TestClient_ReadWriteTransaction_SessionNotFoundOnExecuteBatchUpdate(t *test
 	}
 }
 
+func TestClient_ReadWriteTransaction_WhenLongRunningExecuteBatchUpdate_TakeNoAction(t *testing.T) {
+	t.Parallel()
+	server, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{
+		SessionPoolConfig: SessionPoolConfig{
+			MinOpened: 1,
+			MaxOpened: 1,
+			InactiveTransactionRemovalOptions: InactiveTransactionRemovalOptions{
+				ActionOnInactiveTransaction: WarnAndClose,
+			},
+		},
+	})
+	defer teardown()
+	server.TestSpanner.PutExecutionTime(
+		MethodExecuteBatchDml,
+		SimulatedExecutionTime{Errors: []error{newSessionNotFoundError("projects/p/instances/i/databases/d/sessions/s")}},
+	)
+	ctx := context.Background()
+	p := client.idleSessions
+	var attempts int
+	_, err := client.ReadWriteTransaction(ctx, func(ctx context.Context, tx *ReadWriteTransaction) error {
+		attempts++
+		if attempts == 2 {
+			// Simulate the session to be long-running. The background task should not clean up this long-running session.
+			tx.sh.mu.Lock()
+			tx.sh.checkoutTime = time.Now().Add(-time.Hour)
+			if g, w := tx.sh.eligibleForLongRunning, true; g != w {
+				tx.sh.mu.Unlock()
+				return status.Errorf(codes.FailedPrecondition, "isLongRunningTransaction value mismatch\nGot: %v\nWant: %v", g, w)
+			}
+			tx.sh.mu.Unlock()
+
+			// force run task to clean up unexpected long-running sessions
+			p.removeLongRunningSessions()
+		}
+		rowCounts, err := tx.BatchUpdate(ctx, []Statement{NewStatement(UpdateBarSetFoo)})
+		if err != nil {
+			return err
+		}
+		if g, w := len(rowCounts), 1; g != w {
+			return status.Errorf(codes.FailedPrecondition, "Row counts length mismatch\nGot: %v\nWant: %v", g, w)
+		}
+		if g, w := rowCounts[0], int64(UpdateBarSetFooRowCount); g != w {
+			return status.Errorf(codes.FailedPrecondition, "Row count mismatch\nGot: %v\nWant: %v", g, w)
+		}
+		return nil
+	})
+	if err != nil {
+		t.Fatal(err)
+	}
+	if g, w := attempts, 2; g != w {
+		t.Fatalf("number of attempts mismatch:\nGot%d\nWant:%d", g, w)
+	}
+	p.InactiveTransactionRemovalOptions.mu.Lock()
+	defer p.InactiveTransactionRemovalOptions.mu.Unlock()
+	if g, w := p.numOfLeakedSessionsRemoved, uint64(0); g != w {
+		t.Fatalf("Number of leaked sessions removed mismatch\nGot: %d\nWant: %d\n", g, w)
+	}
+}
+
 func TestClient_ReadWriteTransaction_Query_QueryOptions(t *testing.T) {
 	for _, tt := range queryOptionsTestCases() {
 		t.Run(tt.name, func(t *testing.T) {
@@ -3912,6 +4091,44 @@ func TestClient_PDML_Priority(t *testing.T) {
 	}
 }
 
+func TestClient_WhenLongRunningPartitionedUpdateRequest_TakeNoAction(t *testing.T) {
+	t.Parallel()
+	ctx := context.Background()
+	server, client, teardown := setupMockedTestServerWithConfig(t, ClientConfig{
+		SessionPoolConfig: SessionPoolConfig{
+			MinOpened:                 1,
+			MaxOpened:                 1,
+			healthCheckSampleInterval: 10 * time.Millisecond, // maintainer runs every 10ms
+			InactiveTransactionRemovalOptions: InactiveTransactionRemovalOptions{
+				ActionOnInactiveTransaction: WarnAndClose,
+				executionFrequency:          15 * time.Millisecond, // check long-running sessions every 15ms
+			},
+		},
+	})
+	defer teardown()
+	// delay the rpc by 30ms. The background task runs to clean long-running sessions.
+	server.TestSpanner.PutExecutionTime(MethodExecuteSql,
+		SimulatedExecutionTime{
+			MinimumExecutionTime: 30 * time.Millisecond,
+		})
+
+	stmt := NewStatement(UpdateBarSetFoo)
+	// This transaction is eligible to be long-running, so the background task should not clean its session.
+	rowCount, err := client.PartitionedUpdate(ctx, stmt)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if g, w := rowCount, int64(UpdateBarSetFooRowCount); g != w {
+		t.Errorf("Row count mismatch\nGot: %v\nWant: %v", g, w)
+	}
+	p := client.idleSessions
+	p.InactiveTransactionRemovalOptions.mu.Lock()
+	defer p.InactiveTransactionRemovalOptions.mu.Unlock()
+	if g, w := p.numOfLeakedSessionsRemoved, uint64(0); g != w {
+		t.Fatalf("Number of leaked sessions removed mismatch\nGot: %d\nWant: %d\n", g, w)
+	}
+}
+
 func TestClient_Apply_Priority(t *testing.T) {
 	t.Parallel()
 

spanner/pdml.go

@@ -53,6 +53,10 @@ func (c *Client) partitionedUpdate(ctx context.Context, statement Statement, opt
 	}
 
 	sh, err := c.idleSessions.take(ctx)
+	// Mark isLongRunningTransaction to true, as the session in case of partitioned dml can be long-running
+	sh.mu.Lock()
+	sh.eligibleForLongRunning = true
+	sh.mu.Unlock()
 	if err != nil {
 		return 0, ToSpannerError(err)
 	}