retry_transaction.go

package main

func main() {
    ctx := context.Background()
    db := &db.Client{}

    tx, err := db.Begin()
    if err != nil {
        return err
    }

	// It's recommended to set deadlines on RPCs and around retrying. This is
    // also usually preferred over setting some fixed number of retries: one
    // advantage this has is that backoff settings can be changed independently
    // of the deadline, whereas with a fixed number of retries the deadline
    // would be a constantly-shifting goalpost.
    ctxWithTimeout, cancel := context.WithDeadline(ctx, time.Now().Add(someRPCTimeout))
    defer cancel()

	var someRowsImQuerying []db.Row

	if err := performWithTransactionalRetry(db, func(ctx context.Context, tx *db.Transaction) error {
		if _, err := tx.Update(ctx, "..."); err != nil {
			return err
		}

		if _, err := tx.Create(ctx, "..."); err != nil {
			return err
		}

		if _, err := tx.Insert(ctx, "..."); err != nil {
			return err
		}

		// (probably not how you do reads, but you get the idea)
		if rows, err := tx.Read(ctx, "..."); err != nil {
			someRowsImQuerying = append(someRowsImQuerying, rows...)
		}

		// or, however you do queries with tx

		return
	}); err != nil {
		panic(err)
	}

    resp, err := performSomeRPCWithRetry(ctxWithTimeout)
    if err != nil {
        // TODO: handle err
    }
    _ = resp // TODO: use resp if err is nil
}

func performWithTransactionalRetry(ctx context.Context, db *db.Client, fn func(context.Context, *db.Transaction) error) error {
	// You could also add a max retries here, but usually deadline is better,
	// and that's already on your ctx.
	for {
		tx, err := db.Begin()
		if err != nil {
			return err
		}

		if err := fn(ctx, tx); err != nil {
			if status.Code(err) == ABORTED || strings.Contains(err, "some retryable error") {
				if err := tx.Rollback(); err != nil {
					return err
				}
				continue
			}
			return err
		}

		if err := tx.Commit(); err != nil {
			return err
		}

		break
	}
}