spanner: get read timestamp early

For single-use transactions, the read timestamp wasn't available until after Stop was called on the iterator. Now it is available after the first row is read, as advertised. Fixes Google-internal bug 38008348. Change-Id: Id5ce338e3c610d889003c0d9f9b9a1d4b4de468f Reviewed-on: https://code-review.googlesource.com/13370 Reviewed-by: kokoro <[email protected]> Reviewed-by: Vikas Kedia <[email protected]>
cuecal · Jun 1, 2017 · 3d69fa2 · 3d69fa2
1 parent 0a3a2cd
commit 3d69fa2
Show file tree

Hide file tree

Showing 4 changed files with 97 additions and 25 deletions.
diff --git a/spanner/read.go b/spanner/read.go
@@ -45,24 +45,26 @@ func errEarlyReadEnd() error {
 
 // stream is the internal fault tolerant method for streaming data from
 // Cloud Spanner.
-func stream(ctx context.Context, rpc func(ct context.Context, resumeToken []byte) (streamingReceiver, error), release func(time.Time, error)) *RowIterator {
+func stream(ctx context.Context, rpc func(ct context.Context, resumeToken []byte) (streamingReceiver, error), setTimestamp func(time.Time), release func(error)) *RowIterator {
 	ctx, cancel := context.WithCancel(ctx)
 	return &RowIterator{
-		streamd: newResumableStreamDecoder(ctx, rpc),
-		rowd:    &partialResultSetDecoder{},
-		release: release,
-		cancel:  cancel,
+		streamd:      newResumableStreamDecoder(ctx, rpc),
+		rowd:         &partialResultSetDecoder{},
+		setTimestamp: setTimestamp,
+		release:      release,
+		cancel:       cancel,
 	}
 }
 
 // RowIterator is an iterator over Rows.
 type RowIterator struct {
-	streamd *resumableStreamDecoder
-	rowd    *partialResultSetDecoder
-	release func(time.Time, error)
-	cancel  func()
-	err     error
-	rows    []*Row
+	streamd      *resumableStreamDecoder
+	rowd         *partialResultSetDecoder
+	setTimestamp func(time.Time)
+	release      func(error)
+	cancel       func()
+	err          error
+	rows         []*Row
 }
 
 // Next returns the next result. Its second return value is iterator.Done if
@@ -77,6 +79,10 @@ func (r *RowIterator) Next() (*Row, error) {
 		if r.err != nil {
 			return nil, r.err
 		}
+		if !r.rowd.ts.IsZero() && r.setTimestamp != nil {
+			r.setTimestamp(r.rowd.ts)
+			r.setTimestamp = nil
+		}
 	}
 	if len(r.rows) > 0 {
 		row := r.rows[0]
@@ -123,7 +129,7 @@ func (r *RowIterator) Stop() {
 		r.cancel()
 	}
 	if r.release != nil {
-		r.release(r.rowd.ts, r.err)
+		r.release(r.err)
 		if r.err == nil {
 			r.err = spannerErrorf(codes.FailedPrecondition, "Next called after Stop")
 		}

diff --git a/spanner/read_test.go b/spanner/read_test.go
@@ -1317,7 +1317,8 @@ func TestResumeToken(t *testing.T) {
 				sr.rpcReceiver = r
 				return sr, err
 			},
-			func(time.Time, error) {})
+			nil,
+			func(error) {})
 		defer iter.Stop()
 		for {
 			var row *Row
@@ -1515,7 +1516,8 @@ func TestGrpcReconnect(t *testing.T) {
 				})
 
 			},
-			func(time.Time, error) {})
+			nil,
+			func(error) {})
 		defer iter.Stop()
 		for {
 			_, err = iter.Next()
@@ -1585,7 +1587,8 @@ func TestCancelTimeout(t *testing.T) {
 					ResumeToken: resumeToken,
 				})
 			},
-			func(time.Time, error) {})
+			nil,
+			func(error) {})
 		defer iter.Stop()
 		for {
 			_, err = iter.Next()
@@ -1618,7 +1621,8 @@ func TestCancelTimeout(t *testing.T) {
 					ResumeToken: resumeToken,
 				})
 			},
-			func(time.Time, error) {})
+			nil,
+			func(error) {})
 		defer iter.Stop()
 		for {
 			_, err = iter.Next()
@@ -1667,7 +1671,8 @@ func TestRowIteratorDo(t *testing.T) {
 				ResumeToken: resumeToken,
 			})
 		},
-		func(time.Time, error) {})
+		nil,
+		func(error) {})
 	err = iter.Do(func(r *Row) error { nRows++; return nil })
 	if err != nil {
 		t.Errorf("Using Do: %v", err)
@@ -1702,7 +1707,8 @@ func TestIteratorStopEarly(t *testing.T) {
 				ResumeToken: resumeToken,
 			})
 		},
-		func(time.Time, error) {})
+		nil,
+		func(error) {})
 	_, err = iter.Next()
 	if err != nil {
 		t.Fatalf("before Stop: %v", err)

diff --git a/spanner/spanner_test.go b/spanner/spanner_test.go
@@ -814,6 +814,57 @@ func compareRows(iter *RowIterator, wantNums []int) (string, bool) {
 	return "", true
 }
 
+func TestEarlyTimestamp(t *testing.T) {
+	// Test that we can get the timestamp from a read-only transaction as
+	// soon as we have read at least one row.
+	ctx, cancel := context.WithTimeout(context.Background(), 20*time.Second)
+	defer cancel()
+	// Set up testing environment.
+	if err := prepare(ctx, t, readDBStatements); err != nil {
+		// If prepare() fails, tear down whatever that's already up.
+		tearDown(ctx, t)
+		t.Fatalf("cannot set up testing environment: %v", err)
+	}
+	// After all tests, tear down testing environment.
+	defer tearDown(ctx, t)
+
+	var ms []*Mutation
+	for i := 0; i < 3; i++ {
+		ms = append(ms, InsertOrUpdate(testTable,
+			testTableColumns,
+			[]interface{}{fmt.Sprintf("k%d", i), fmt.Sprintf("v%d", i)}))
+	}
+	if _, err := client.Apply(ctx, ms, ApplyAtLeastOnce()); err != nil {
+		t.Fatal(err)
+	}
+
+	txn := client.Single()
+	iter := txn.Read(ctx, testTable, AllKeys(), testTableColumns)
+	defer iter.Stop()
+	// In  single-use transaction, we should get an error before reading anything.
+	if _, err := txn.Timestamp(); err == nil {
+		t.Error("wanted error, got nil")
+	}
+	// After reading one row, the timestamp should be available.
+	_, err := iter.Next()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if _, err := txn.Timestamp(); err != nil {
+		t.Errorf("got %v, want nil", err)
+	}
+
+	txn = client.ReadOnlyTransaction()
+	defer txn.Close()
+	iter = txn.Read(ctx, testTable, AllKeys(), testTableColumns)
+	defer iter.Stop()
+	// In an ordinary read-only transaction, the timestamp should be
+	// available immediately.
+	if _, err := txn.Timestamp(); err != nil {
+		t.Errorf("got %v, want nil", err)
+	}
+}
+
 func TestNestedTransaction(t *testing.T) {
 	// You cannot use a transaction from inside a read-write transaction.
 	ctx := context.Background()

diff --git a/spanner/transaction.go b/spanner/transaction.go
@@ -36,8 +36,10 @@ type transactionID []byte
 type txReadEnv interface {
 	// acquire returns a read-transaction environment that can be used to perform a transactional read.
 	acquire(ctx context.Context) (*sessionHandle, *sppb.TransactionSelector, error)
-	// release should be called at the end of every transactional read to deal with session recycling and read timestamp recording.
-	release(time.Time, error)
+	// sets the transaction's read timestamp
+	setTimestamp(time.Time)
+	// release should be called at the end of every transactional read to deal with session recycling.
+	release(error)
 }
 
 // txReadOnly contains methods for doing transactional reads.
@@ -97,6 +99,7 @@ func (t *txReadOnly) ReadUsingIndex(ctx context.Context, table, index string, ke
 					ResumeToken: resumeToken,
 				})
 		},
+		t.setTimestamp,
 		t.release,
 	)
 }
@@ -155,6 +158,7 @@ func (t *txReadOnly) Query(ctx context.Context, statement Statement) *RowIterato
 			req.ResumeToken = resumeToken
 			return client.ExecuteStreamingSql(ctx, req)
 		},
+		t.setTimestamp,
 		t.release)
 }
 
@@ -409,12 +413,17 @@ func (t *ReadOnlyTransaction) acquireMultiUse(ctx context.Context) (*sessionHand
 	}
 }
 
-// release implements txReadEnv.release.
-func (t *ReadOnlyTransaction) release(rts time.Time, err error) {
+func (t *ReadOnlyTransaction) setTimestamp(ts time.Time) {
 	t.mu.Lock()
-	if t.singleUse && !rts.IsZero() {
-		t.rts = rts
+	defer t.mu.Unlock()
+	if t.rts.IsZero() {
+		t.rts = ts
 	}
+}
+
+// release implements txReadEnv.release.
+func (t *ReadOnlyTransaction) release(err error) {
+	t.mu.Lock()
 	sh := t.sh
 	t.mu.Unlock()
 	if sh != nil { // sh could be nil if t.acquire() fails.
@@ -599,7 +608,7 @@ func (t *ReadWriteTransaction) acquire(ctx context.Context) (*sessionHandle, *sp
 }
 
 // release implements txReadEnv.release.
-func (t *ReadWriteTransaction) release(_ time.Time, err error) {
+func (t *ReadWriteTransaction) release(err error) {
 	t.mu.Lock()
 	sh := t.sh
 	t.mu.Unlock()