spanner: Race condition in (*BatchReadOnlyTransaction).Execute()

[lhecker]

Client

Spanner

Environment

Any

Code

The following code counts the total number of rows using PartitionRead().
It'll exhibit the race condition by showing varying total counts at the end of each run.
Of course this only works if your table has more than one partition.

package main

import (
	"context"
	"fmt"
	"sync"
	"sync/atomic"

	"cloud.google.com/go/spanner"
)

func main() {
	client, err := spanner.NewClient(context.Background(), "projects/<project>/instances/<instance>/databases/<database>")
	if err != nil {
		fmt.Println(err)
		return
	}

	txn, err := client.BatchReadOnlyTransaction(context.Background(), spanner.StrongRead())
	if err != nil {
		fmt.Println(err)
		return
	}
	defer txn.Cleanup(context.Background())

	partitions, err := txn.PartitionRead(context.Background(), "<table>", spanner.AllKeys(), []string{"<column>"}, spanner.PartitionOptions{})
	if err != nil {
		fmt.Println(err)
		return
	}

	wg := &sync.WaitGroup{}
	var total uint64

	for _, p := range partitions {
		p := p

		go func() {
			iter := txn.Execute(context.Background(), p)
			defer iter.Stop()

			var count uint64
			err = iter.Do(func(row *spanner.Row) error {
				count++
				return nil
			})
			if err != nil {
				return
			}

			atomic.AddUint64(&total, count)
		}()
	}

	wg.Wait()
	fmt.Println(total)
}

Expected behavior

PartitionRead() reads all rows.

Actual behavior

A race condition. 😢

Cause

Each partition points to the same rreq/qreq.
For instance for rreq line 137 in the following code section:

google-cloud-go/spanner/batch.go

Lines 124 to 139 in a1a1bde

	// Prepare ReadRequest.
	req := &sppb.ReadRequest{
	Session: sid,
	Transaction: ts,
	Table: table,
	Index: index,
	Columns: columns,
	KeySet: kset,
	}
	// Generate partitions.
	for _, p := range resp.GetPartitions() {
	partitions = append(partitions, &Partition{
	pt: p.PartitionToken,
	rreq: req,
	})
	}

During the call in Execution() the rreq/qreq fields are then mutated without locks of any kind.
For instance for rreq line 264 in the following code section:

google-cloud-go/spanner/batch.go

Lines 262 to 275 in a1a1bde

	// Read or query partition.
	if p.rreq != nil {
	p.rreq.PartitionToken = p.pt
	rpc = func(ctx context.Context, resumeToken []byte) (streamingReceiver, error) {
	p.rreq.ResumeToken = resumeToken
	return client.StreamingRead(ctx, p.rreq)
	}
	} else {
	p.qreq.PartitionToken = p.pt
	rpc = func(ctx context.Context, resumeToken []byte) (streamingReceiver, error) {
	p.qreq.ResumeToken = resumeToken
	return client.ExecuteStreamingSql(ctx, p.qreq)
	}
	}

The documentation of that function also doesn't mention it's not safe for concurrent use, nor that it mutates it's argument.

I suppose the struct field pt exists because this was an attempt to reduce the data size after a call to MarshalBinary()?
In that case I'd like to kindly suggest to simply copy the rreq/qreq fields within the Execution() method and mutate the copy only. 🙂

Additional context

Intermediate fix:

data, err := p.MarshalBinary()
if err != nil {
	return
}

p := &spanner.Partition{}
err = p.UnmarshalBinary(data)
if err != nil {
	return
}

👉 Should I submit a fix?
I already did so in the past, but last time you told me to open an issue beforehand. 😄

[olavloite]

@lhecker Thanks, that's a great catch. And thank you for the very elaborate report and analysis.

There's a fix coming here: https://code-review.googlesource.com/c/gocloud/+/54190