| 1 | package console |
| 2 | |
| 3 | import ( |
| 4 | "bytes" |
| 5 | "context" |
| 6 | "database/sql" |
| 7 | "fmt" |
| 8 | "runtime" |
| 9 | "strconv" |
| 10 | ) |
| 11 | |
| 12 | // errOuterStoreInView is a call on the outer, captured Store from inside the |
| 13 | // callback withClaimSQL is running (review m-5). The view holds the pool's |
| 14 | // only connection, so the call could only wait for its own caller; it is |
| 15 | // refused at once instead, before any remote check. |
| 16 | var errOuterStoreInView = fmt.Errorf("%w: the outer Store was used inside withClaimSQL; use the view it passes", ErrClaimSQLScope) |
| 17 | |
| 18 | // goroutineID is the running goroutine's number. It is used only to tell |
| 19 | // whether the caller is the goroutine that holds the bound view; a call from |
| 20 | // any other goroutine waits for the connection as before. |
| 21 | func goroutineID() uint64 { |
| 22 | var buf [64]byte |
| 23 | b := bytes.TrimPrefix(buf[:runtime.Stack(buf[:], false)], []byte("goroutine ")) |
| 24 | if i := bytes.IndexByte(b, ' '); i > 0 { |
| 25 | b = b[:i] |
| 26 | } |
| 27 | id, _ := strconv.ParseUint(string(b), 10, 64) |
| 28 | return id |
| 29 | } |
| 30 | |
| 31 | // insideView reports whether the calling goroutine holds the bound view. |
| 32 | func (g *claimGate) insideView() bool { |
| 33 | h := g.viewG.Load() |
| 34 | return h != 0 && h == goroutineID() |
| 35 | } |
| 36 | |
| 37 | // outerInsideView reports a call on the outer Store from the goroutine that |
| 38 | // holds its bound view. A goroutine fn starts is not caught: it waits. |
| 39 | func (s *Store) outerInsideView() bool { |
| 40 | return s.gate != nil && s.bound == nil && s.gate.insideView() |
| 41 | } |
| 42 | |
| 43 | // claimOuterDB is the gated Store's own handle. It is the pool, except that |
| 44 | // the goroutine holding the bound view is refused instead of waiting on the |
| 45 | // connection it already holds. |
| 46 | type claimOuterDB struct { |
| 47 | *sql.DB |
| 48 | g *claimGate |
| 49 | } |
| 50 | |
| 51 | var _ sqlHandle = claimOuterDB{} |
| 52 | |
| 53 | // refusedContext is already done with errOuterStoreInView, so database/sql |
| 54 | // refuses before it waits for a connection. It is how QueryRow, which has no |
| 55 | // error return, carries the refusal to Scan. |
| 56 | type refusedContext struct{ context.Context } |
| 57 | |
| 58 | var closedDone = func() chan struct{} { c := make(chan struct{}); close(c); return c }() |
| 59 | |
| 60 | func (refusedContext) Done() <-chan struct{} { return closedDone } |
| 61 | func (refusedContext) Err() error { return errOuterStoreInView } |
| 62 | |
| 63 | func (o claimOuterDB) guard(ctx context.Context) context.Context { |
| 64 | if o.g.insideView() { |
| 65 | return refusedContext{ctx} |
| 66 | } |
| 67 | return ctx |
| 68 | } |
| 69 | |
| 70 | func (o claimOuterDB) Exec(query string, args ...any) (sql.Result, error) { |
| 71 | return o.ExecContext(context.Background(), query, args...) |
| 72 | } |
| 73 | func (o claimOuterDB) ExecContext(ctx context.Context, query string, args ...any) (sql.Result, error) { |
| 74 | return o.DB.ExecContext(o.guard(ctx), query, args...) |
| 75 | } |
| 76 | func (o claimOuterDB) Query(query string, args ...any) (*sql.Rows, error) { |
| 77 | return o.QueryContext(context.Background(), query, args...) |
| 78 | } |
| 79 | func (o claimOuterDB) QueryContext(ctx context.Context, query string, args ...any) (*sql.Rows, error) { |
| 80 | return o.DB.QueryContext(o.guard(ctx), query, args...) |
| 81 | } |
| 82 | func (o claimOuterDB) QueryRow(query string, args ...any) *sql.Row { |
| 83 | return o.QueryRowContext(context.Background(), query, args...) |
| 84 | } |
| 85 | func (o claimOuterDB) QueryRowContext(ctx context.Context, query string, args ...any) *sql.Row { |
| 86 | return o.DB.QueryRowContext(o.guard(ctx), query, args...) |
| 87 | } |
| 88 | func (o claimOuterDB) Begin() (*sql.Tx, error) { return o.BeginTx(context.Background(), nil) } |
| 89 | func (o claimOuterDB) BeginTx(ctx context.Context, opts *sql.TxOptions) (*sql.Tx, error) { |
| 90 | return o.DB.BeginTx(o.guard(ctx), opts) |
| 91 | } |
| 92 | func (o claimOuterDB) Prepare(query string) (*sql.Stmt, error) { |
| 93 | return o.PrepareContext(context.Background(), query) |
| 94 | } |
| 95 | func (o claimOuterDB) PrepareContext(ctx context.Context, query string) (*sql.Stmt, error) { |
| 96 | return o.DB.PrepareContext(o.guard(ctx), query) |
| 97 | } |
| 98 | |
amadan