go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/gae/filter/txndefer/filter.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 txndefer implements a filter that calls best-effort callbacks on
// successful transaction commits.
//
// Useful when an activity inside a transaction has some best-effort follow up
// that should be done once the transaction has successfully landed.
package txndefer

import (
	"context"
	"sync"

	ds "go.chromium.org/luci/gae/service/datastore"
)

// FilterRDS installs the datastore filter into the context.
func FilterRDS(ctx context.Context) context.Context {
	return ds.AddRawFilters(ctx, func(_ context.Context, inner ds.RawInterface) ds.RawInterface {
		return filteredDS{inner}
	})
}

// Defer schedules `cb` for execution when the current 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 RunInTransaction 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
// RunInTransaction so that they inherit the deadline of the entire transaction.
//
// Panics if the given context is not transactional or there's no txndefer
// filter installed.
func Defer(ctx context.Context, cb func(context.Context)) {
	state, _ := ctx.Value(&ctxKey).(*txnState)
	if state == nil {
		panic("not a transactional context or no txndefer filter installed")
	}
	state.deferCB(cb)
}

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

var ctxKey = "txndefer.txnState"

type txnState struct {
	ctx context.Context // the original transaction context
	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() {
	// Note: execCBs happens after RunInTransaction 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.
	if len(s.cbs) != 0 {
		ctx := ds.WithoutTransaction(s.ctx)
		for i := len(s.cbs) - 1; i >= 0; i-- {
			s.cbs[i](ctx)
		}
	}
}

type filteredDS struct {
	ds.RawInterface
}

func (fds filteredDS) RunInTransaction(f func(ctx context.Context) error, opts *ds.TransactionOptions) error {
	var state *txnState
	err := fds.RawInterface.RunInTransaction(func(ctx context.Context) error {
		state = &txnState{ctx: ctx}
		return f(context.WithValue(ctx, &ctxKey, state))
	}, opts)
	if err == nil {
		state.execCBs()
	}
	return err
}