go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/server/span/client.go

// Copyright 2020 The LUCI Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package span

import (
	"context"
	"sync"
	"time"

	"cloud.google.com/go/spanner"
	sppb "cloud.google.com/go/spanner/apiv1/spannerpb"
	"google.golang.org/grpc/codes"

	"go.chromium.org/luci/common/errors"
)

// Transaction is a common interface of spanner.ReadOnlyTransaction and
// spanner.ReadWriteTransaction.
type Transaction interface {
	// ReadRow reads a single row from the database.
	ReadRow(ctx context.Context, table string, key spanner.Key, columns []string) (*spanner.Row, error)

	// ReadRowWithOptions reads a single row from the database, and allows customizing options.
	ReadRowWithOptions(ctx context.Context, table string, key spanner.Key, columns []string, opts *spanner.ReadOptions) (*spanner.Row, error)

	// Read reads multiple rows from the database.
	Read(ctx context.Context, table string, keys spanner.KeySet, columns []string) *spanner.RowIterator

	// ReadWithOptions reads multiple rows from the database, and allows
	// customizing options.
	ReadWithOptions(ctx context.Context, table string, keys spanner.KeySet, columns []string, opts *spanner.ReadOptions) *spanner.RowIterator

	// Query reads multiple rows returned by SQL statement.
	Query(ctx context.Context, statement spanner.Statement) *spanner.RowIterator

	// QueryWithOptions reads multiple rows returned by SQL statement, and allows
	// customizing options.
	QueryWithOptions(ctx context.Context, statement spanner.Statement, opts spanner.QueryOptions) *spanner.RowIterator
}

// UseClient installs a Spanner client into the context.
//
// Primarily used by the module initialization code. May be useful in tests as
// well.
func UseClient(ctx context.Context, client *spanner.Client) context.Context {
	return context.WithValue(ctx, &clientContextKey, client)
}

// Apply applies a list of mutations atomically to the database.
//
// Panics if called from inside a read-write transaction. Use BufferWrite to
// apply mutations there.
func Apply(ctx context.Context, ms []*spanner.Mutation, opts ...spanner.ApplyOption) (commitTimestamp time.Time, err error) {
	panicOnNestedRW(ctx)
	if ctxOpts, ok := ctx.Value(&requestOptionsContextKey).(*RequestOptions); ok {
		opts = append([]spanner.ApplyOption{spanner.Priority(ctxOpts.Priority)}, opts...)
	}
	return client(ctx).Apply(ctx, ms, opts...)
}

// PartitionedUpdate executes a DML statement in parallel across the database,
// using separate, internal transactions that commit independently.
//
// Panics if called from inside a read-write transaction.
func PartitionedUpdate(ctx context.Context, st spanner.Statement) (count int64, err error) {
	panicOnNestedRW(ctx)
	return client(ctx).PartitionedUpdateWithOptions(ctx, st, queryOptionsFromContext(ctx))
}

// Single returns a derived context with a single-use read-only transaction.
//
// It provides a read-only snapshot transaction optimized for the case where
// only a single read or query is needed. This is more efficient than using
// ReadOnlyTransaction() for a single read or query.
//
// The transaction object can be obtained through RO(ctx) or Txn(ctx). It is
// also transparently used by ReadRow, Read, Query, etc. wrappers.
//
// Panics if `ctx` already holds a transaction (either read-write or read-only).
func Single(ctx context.Context) context.Context {
	panicOnNestedRO(ctx)
	return setTxnState(ctx, &txnState{ro: client(ctx).Single()})
}

// ReadOnlyTransaction returns a derived context with a read-only transaction.
//
// It can be used for multiple reads from the database. To avoid leaking
// resources on the server this context *must* be canceled as soon as all reads
// are done.
//
// The transaction object can be obtained through RO(ctx) or Txn(ctx). It is
// also transparently used by ReadRow, Read, Query, etc. wrappers.
//
// Panics if `ctx` already holds a transaction (either read-write or read-only).
func ReadOnlyTransaction(ctx context.Context) (context.Context, context.CancelFunc) {
	panicOnNestedRO(ctx)
	txn := client(ctx).ReadOnlyTransaction()
	ctx, cancel := context.WithCancel(setTxnState(ctx, &txnState{ro: txn}))
	return ctx, func() { cancel(); txn.Close() }
}

// ReadWriteTransaction executes a read-write transaction, with retries as
// necessary.
//
// The callback may be called multiple times if Spanner client decides to retry
// the transaction. In particular this happens if the callback returns (perhaps
// wrapped) ABORTED error. This error is returned by Spanner client methods if
// they encounter a stale transaction.
//
// See https://godoc.org/cloud.google.com/go/spanner#ReadWriteTransaction for
// more details.
//
// The callback can access the transaction object via RW(ctx) or Txn(ctx). It is
// also transparently used by ReadRow, Read, Query, BufferWrite, etc. wrappers.
//
// Panics if `ctx` already holds a read-write transaction. Starting a read-write
// transaction from a read-only transaction is OK though, but beware that they
// are completely separate unrelated transactions.
func ReadWriteTransaction(ctx context.Context, f func(ctx context.Context) error) (commitTimestamp time.Time, err error) {
	panicOnNestedRW(ctx)

	var state *txnState

	cts, err := client(ctx).ReadWriteTransaction(ctx, func(ctx context.Context, rw *spanner.ReadWriteTransaction) error {
		state = &txnState{rw: rw}
		err := f(setTxnState(ctx, state))
		if unwrapped := errors.Unwrap(err); spanner.ErrCode(unwrapped) == codes.Aborted {
			err = unwrapped
		}
		return err
	})

	if err == nil {
		state.execCBs(ctx)
	}

	return cts, err
}

// RO returns the current read-only transaction in the context or nil if it's
// not a read-only transactional context.
func RO(ctx context.Context) *spanner.ReadOnlyTransaction {
	if s := getTxnState(ctx); s != nil {
		return s.ro
	}
	return nil
}

// MustRO is like RO except it panics if `ctx` is not read-only transactional.
func MustRO(ctx context.Context) *spanner.ReadOnlyTransaction {
	if ro := RO(ctx); ro != nil {
		return ro
	}
	panic("not a read-only Spanner transactional context")
}

// RW returns the current read-write transaction in the context or nil if it's
// not a read-write transactional context.
func RW(ctx context.Context) *spanner.ReadWriteTransaction {
	if s := getTxnState(ctx); s != nil {
		return s.rw
	}
	return nil
}

// MustRW is like RW except it panics if `ctx` is not read-write transactional.
func MustRW(ctx context.Context) *spanner.ReadWriteTransaction {
	if rw := RW(ctx); rw != nil {
		return rw
	}
	panic("not a read-write Spanner transactional context")
}

// Txn returns an interface that can be used to read data in the current
// read-only or read-write transaction.
//
// Returns nil if `ctx` is not a transactional context.
func Txn(ctx context.Context) Transaction {
	switch s := getTxnState(ctx); {
	case s == nil:
		return nil
	case s.ro != nil:
		return s.ro
	default:
		return s.rw
	}
}

// MustTxn is like Txn except it panics if `ctx` is not transactional.
func MustTxn(ctx context.Context) Transaction {
	if txn := Txn(ctx); txn != nil {
		return txn
	}
	panic("not a transactional Spanner context")
}

// WithoutTxn returns a copy of the context without the transaction in it.
//
// This can be used to spawn separate independent transactions from within
// a transaction.
func WithoutTxn(ctx context.Context) context.Context {
	if getTxnState(ctx) == nil {
		return ctx
	}
	return setTxnState(ctx, nil)
}

// Defer schedules `cb` for execution when the current read-write transaction
// successfully lands.
//
// Intended for a best-effort non-transactional follow up to a successful
// transaction. Note that in presence of failures there's no guarantee the
// callback will be called. For example, the callback won't ever be called if
// the process crashes right after landing the transaction. Or if the
// transaction really landed, but ReadWriteTransaction finished with
// "deadline exceeded" (or some similar) error.
//
// Callbacks are executed sequentially in the reverse order they were deferred.
// They receive the non-transactional version of the context initially passed to
// ReadWriteTransaction.
//
// Panics if the given context is not transactional.
func Defer(ctx context.Context, cb func(context.Context)) {
	state := getTxnState(ctx)
	if state == nil || state.rw == nil {
		panic("not a read-write Spanner transactional context")
	}
	state.deferCB(cb)
}

// ReadRow reads a single row from the database.
//
// It is a shortcut for MustTxn(ctx).ReadRow(ctx, ...). Panics if the context
// is not transactional.
func ReadRow(ctx context.Context, table string, key spanner.Key, columns []string) (*spanner.Row, error) {
	return MustTxn(ctx).ReadRowWithOptions(ctx, table, key, columns, readOptionsFromContext(ctx))
}

// ReadRowWithOptions reads a single row from the database, and allows customizing
// options.
//
// It is a shortcut for MustTxn(ctx).ReadRowWithOptions(ctx, ...). Panics if
// the context is not transactional.
//
// It does not use the default RequestOptions from the ctx.  Use opts to pass
// these explicitly.
func ReadRowWithOptions(ctx context.Context, table string, key spanner.Key, columns []string, opts *spanner.ReadOptions) (*spanner.Row, error) {
	return MustTxn(ctx).ReadRowWithOptions(ctx, table, key, columns, opts)
}

// Read reads multiple rows from the database.
//
// It is a shortcut for MustTxn(ctx).Read(ctx, ...). Panics if the context
// is not transactional.
func Read(ctx context.Context, table string, keys spanner.KeySet, columns []string) *spanner.RowIterator {
	return MustTxn(ctx).ReadWithOptions(ctx, table, keys, columns, readOptionsFromContext(ctx))
}

// ReadWithOptions reads multiple rows from the database, and allows customizing
// options.
//
// It is a shortcut for MustTxn(ctx).ReadWithOptions(ctx, ...). Panics if the
// context is not transactional.
//
// It does not use the default RequestOptions from the ctx.  Use opts to pass
// these explicitly.
func ReadWithOptions(ctx context.Context, table string, keys spanner.KeySet, columns []string, opts *spanner.ReadOptions) *spanner.RowIterator {
	return MustTxn(ctx).ReadWithOptions(ctx, table, keys, columns, opts)
}

// Query reads multiple rows returned by SQL statement.
//
// It is a shortcut for MustTxn(ctx).Query(ctx, ...). Panics if the context is
// not transactional.
func Query(ctx context.Context, statement spanner.Statement) *spanner.RowIterator {
	return MustTxn(ctx).QueryWithOptions(ctx, statement, queryOptionsFromContext(ctx))
}

// BufferWrite adds a list of mutations to the set of updates that will be
// applied when the transaction is committed.
//
// It does not actually apply the write until the transaction is committed, so
// the operation does not block. The effects of the write won't be visible to
// any reads (including reads done in the same transaction) until the
// transaction commits.
//
// It is a wrapper over MustRW(ctx).BufferWrite(...). Panics if the context is
// not read-write transactional.
func BufferWrite(ctx context.Context, ms ...*spanner.Mutation) {
	// BufferWrite just appends mutation to an internal buffer. It fails only if
	// the transaction has already landed. We are OK to panic in this case:
	// calling BufferWrite outside of a transaction is a programming error.
	if err := MustRW(ctx).BufferWrite(ms); err != nil {
		panic(err)
	}
}

// Update executes a DML statement against the database. It returns the number
// of affected rows. Update returns an error if the statement is a query.
// However, the query is executed, and any data read will be validated upon
// commit.
//
// It is a shortcut for MustRW(ctx).Update(...). Panics if the context is not
// read-write transactional.
func Update(ctx context.Context, stmt spanner.Statement) (rowCount int64, err error) {
	return MustRW(ctx).UpdateWithOptions(ctx, stmt, queryOptionsFromContext(ctx))
}

// UpdateWithOptions executes a DML statement against the database. It returns
// the number of affected rows. The sql query execution will be optimized based
// on the given query options.
//
// It is a shortcut for MustRW(ctx).Update(...). Panics if the context is not
// read-write transactional.
//
// It does not use the default RequestOptions from the ctx.  Use opts to pass
// these explicitly.
func UpdateWithOptions(ctx context.Context, stmt spanner.Statement, opts spanner.QueryOptions) (rowCount int64, err error) {
	return MustRW(ctx).UpdateWithOptions(ctx, stmt, opts)
}

////////////////////////////////////////////////////////////////////////////////

var (
	clientContextKey         = "go.chromium.org/luci/server/span:client"
	txnContextKey            = "go.chromium.org/luci/server/span:txn"
	requestOptionsContextKey = "go.chromium.org/luci/server/span:requestOptions"
)

// client returns a Spanner client installed in the context.
//
// Panics if it is not there.
//
// Intentionally private to force all callers to go through package's functions
// like ReadWriteTransaction, ReadOnlyTransaction, Single, etc. since they
// generally add additional functionality on top of the raw Spanner client that
// other LUCI packages assume to be present. Using the Spanner client directly
// may violate such assumptions leading to undefined behavior when multiple
// packages are used together.
func client(ctx context.Context) *spanner.Client {
	cl, _ := ctx.Value(&clientContextKey).(*spanner.Client)
	if cl == nil {
		panic("no Spanner client in the context")
	}
	return cl
}

type txnState struct {
	ro *spanner.ReadOnlyTransaction
	rw *spanner.ReadWriteTransaction

	m   sync.Mutex
	cbs []func(context.Context)
}

func (s *txnState) deferCB(cb func(context.Context)) {
	s.m.Lock()
	s.cbs = append(s.cbs, cb)
	s.m.Unlock()
}

func (s *txnState) execCBs(ctx context.Context) {
	// Note: execCBs happens after ReadWriteTransaction has finished. If it
	// spawned any goroutines, they must have been finished already too (calling
	// Defer from a goroutine that outlives a transaction is rightfully a race).
	// Thus all writes to `s.cbs` are finished already and we also passed some
	// synchronization barrier that waited for the goroutines to join. It's fine
	// to avoid locking s.m in this case saving 200ns on hot code path.
	for i := len(s.cbs) - 1; i >= 0; i-- {
		s.cbs[i](ctx)
	}
}

func setTxnState(ctx context.Context, s *txnState) context.Context {
	return context.WithValue(ctx, &txnContextKey, s)
}

func getTxnState(ctx context.Context) *txnState {
	s, _ := ctx.Value(&txnContextKey).(*txnState)
	return s
}

func panicOnNestedRW(ctx context.Context) {
	if RW(ctx) != nil {
		panic("nested Spanner write transactions are not allowed")
	}
}

func panicOnNestedRO(ctx context.Context) {
	if getTxnState(ctx) != nil {
		panic("nested Spanner read transactions are not allowed")
	}
}

// RequestOptions holds options common to many Spanner requests.
//
// Used for setting default request options in the context via
// ModifyRequestOptions.  See also spanner.ReadOptions also
// spanner.QueryOptions.
type RequestOptions struct {
	Tag      string
	Priority sppb.RequestOptions_Priority
}

// ModifyRequestOptions returns a new Context that carries default Spanner
// request options.
//
// These request options will be used by all Spanner operations in this module
// if the underlying operation supports it, except the *WithOptions functions
// (where the caller provides options explicitly).
//
// The cb function will be called with the current defaults from the context,
// to allow for incremental updates.
func ModifyRequestOptions(ctx context.Context, cb func(*RequestOptions)) context.Context {
	var next RequestOptions
	if cur, ok := ctx.Value(&requestOptionsContextKey).(*RequestOptions); ok {
		next = *cur
	}
	cb(&next)
	return context.WithValue(ctx, &requestOptionsContextKey, &next)
}

func queryOptionsFromContext(ctx context.Context) spanner.QueryOptions {
	opts, ok := ctx.Value(&requestOptionsContextKey).(*RequestOptions)
	if !ok {
		return spanner.QueryOptions{}
	}
	return spanner.QueryOptions{
		Priority:   opts.Priority,
		RequestTag: opts.Tag,
	}
}

func readOptionsFromContext(ctx context.Context) *spanner.ReadOptions {
	opts, ok := ctx.Value(&requestOptionsContextKey).(*RequestOptions)
	if !ok {
		return &spanner.ReadOptions{}
	}
	return &spanner.ReadOptions{
		Priority:   opts.Priority,
		RequestTag: opts.Tag,
	}
}