merged from gramsrv upstream
This commit is contained in:
onysd
parent
79c64ee916
commit
21a0856587
651 changed files with 54774 additions and 4590 deletions
408
internal/store/postgres/authkey_session_layer_batch.go
Normal file
408
internal/store/postgres/authkey_session_layer_batch.go
Normal file
|
|
@ -0,0 +1,408 @@
|
|||
package postgres
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"telesrv/internal/store"
|
||||
)
|
||||
|
||||
// AuthKeySessionLayerBatchConfig bounds the synchronous cross-session batch.
|
||||
// A batch never contains the same raw auth-key/session identity twice and a
|
||||
// caller does not return until its batch has committed or failed.
|
||||
type AuthKeySessionLayerBatchConfig struct {
|
||||
MaxSize int
|
||||
MaxWait time.Duration
|
||||
QueueSize int
|
||||
QueryTimeout time.Duration
|
||||
}
|
||||
|
||||
type authKeySessionLayerBatchKey struct {
|
||||
rawAuthKeyID [8]byte
|
||||
sessionID int64
|
||||
}
|
||||
|
||||
type authKeySessionLayerBatchRequest struct {
|
||||
ctx context.Context
|
||||
rawAuthKeyID [8]byte
|
||||
sessionID int64
|
||||
layer int
|
||||
msgID int64
|
||||
expiresAt time.Time
|
||||
result chan authKeySessionLayerBatchResult
|
||||
}
|
||||
|
||||
type authKeySessionLayerBatchResult struct {
|
||||
current store.AuthKeySessionLayer
|
||||
fast bool
|
||||
err error
|
||||
}
|
||||
|
||||
// BatchedAuthKeySessionLayerStore preserves AuthKeySessionLayerStore semantics
|
||||
// while combining contemporaneous same-Layer fast attempts for distinct
|
||||
// sessions into one PostgreSQL statement. A miss is resolved synchronously by
|
||||
// the original full identity transaction before the caller returns.
|
||||
type BatchedAuthKeySessionLayerStore struct {
|
||||
base *AuthKeyStore
|
||||
cfg AuthKeySessionLayerBatchConfig
|
||||
queue chan authKeySessionLayerBatchRequest
|
||||
stop chan struct{}
|
||||
done chan struct{}
|
||||
cancel context.CancelFunc
|
||||
once sync.Once
|
||||
gate sync.RWMutex
|
||||
closed bool
|
||||
}
|
||||
|
||||
func NewBatchedAuthKeySessionLayerStore(
|
||||
base *AuthKeyStore,
|
||||
cfg AuthKeySessionLayerBatchConfig,
|
||||
) (*BatchedAuthKeySessionLayerStore, error) {
|
||||
if base == nil || base.db == nil {
|
||||
return nil, errors.New("initialize auth key session Layer batcher: nil store")
|
||||
}
|
||||
if cfg.MaxSize <= 0 || cfg.MaxSize > 4096 {
|
||||
return nil, fmt.Errorf("initialize auth key session Layer batcher: max size %d outside [1,4096]", cfg.MaxSize)
|
||||
}
|
||||
if cfg.MaxWait <= 0 || cfg.MaxWait > 10*time.Millisecond {
|
||||
return nil, fmt.Errorf("initialize auth key session Layer batcher: max wait %v outside (0,10ms]", cfg.MaxWait)
|
||||
}
|
||||
if cfg.QueueSize < cfg.MaxSize || cfg.QueueSize > 1<<20 {
|
||||
return nil, fmt.Errorf("initialize auth key session Layer batcher: queue size %d outside [%d,%d]", cfg.QueueSize, cfg.MaxSize, 1<<20)
|
||||
}
|
||||
if cfg.QueryTimeout <= 0 || cfg.QueryTimeout > 30*time.Second {
|
||||
return nil, fmt.Errorf("initialize auth key session Layer batcher: query timeout %v outside (0,30s]", cfg.QueryTimeout)
|
||||
}
|
||||
workerCtx, cancel := context.WithCancel(context.Background())
|
||||
s := &BatchedAuthKeySessionLayerStore{
|
||||
base: base, cfg: cfg,
|
||||
queue: make(chan authKeySessionLayerBatchRequest, cfg.QueueSize),
|
||||
stop: make(chan struct{}), done: make(chan struct{}), cancel: cancel,
|
||||
}
|
||||
go s.run(workerCtx)
|
||||
return s, nil
|
||||
}
|
||||
|
||||
func (s *BatchedAuthKeySessionLayerStore) GetSessionLayer(
|
||||
ctx context.Context,
|
||||
rawAuthKeyID [8]byte,
|
||||
sessionID int64,
|
||||
) (store.AuthKeySessionLayer, bool, error) {
|
||||
return s.base.GetSessionLayer(ctx, rawAuthKeyID, sessionID)
|
||||
}
|
||||
|
||||
func (s *BatchedAuthKeySessionLayerStore) AdvanceSessionLayer(
|
||||
ctx context.Context,
|
||||
rawAuthKeyID [8]byte,
|
||||
sessionID int64,
|
||||
layer int,
|
||||
msgID int64,
|
||||
) (store.AuthKeySessionLayer, bool, error) {
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
expiresAt, validMessageID := store.AuthKeySessionLayerExpiry(msgID)
|
||||
if layer <= 0 || !validMessageID {
|
||||
return store.AuthKeySessionLayer{}, false, store.ErrAuthKeySessionLayerInvalid
|
||||
}
|
||||
request := authKeySessionLayerBatchRequest{
|
||||
ctx: ctx, rawAuthKeyID: rawAuthKeyID, sessionID: sessionID,
|
||||
layer: layer, msgID: msgID, expiresAt: expiresAt,
|
||||
result: make(chan authKeySessionLayerBatchResult, 1),
|
||||
}
|
||||
s.gate.RLock()
|
||||
if s.closed {
|
||||
s.gate.RUnlock()
|
||||
return store.AuthKeySessionLayer{}, false, context.Canceled
|
||||
}
|
||||
select {
|
||||
case s.queue <- request:
|
||||
case <-ctx.Done():
|
||||
s.gate.RUnlock()
|
||||
return store.AuthKeySessionLayer{}, false, ctx.Err()
|
||||
}
|
||||
s.gate.RUnlock()
|
||||
|
||||
// Once accepted by the bounded queue, wait for the worker's definitive
|
||||
// commit/error. This prevents a canceled caller from turning the submitted
|
||||
// selector into an unobserved asynchronous best-effort write.
|
||||
result := <-request.result
|
||||
if result.err != nil {
|
||||
return store.AuthKeySessionLayer{}, false, result.err
|
||||
}
|
||||
if result.fast {
|
||||
return result.current, true, nil
|
||||
}
|
||||
if err := ctx.Err(); err != nil {
|
||||
return store.AuthKeySessionLayer{}, false, err
|
||||
}
|
||||
return s.base.advanceSessionLayerFull(ctx, rawAuthKeyID, sessionID, layer, msgID, expiresAt)
|
||||
}
|
||||
|
||||
func (s *BatchedAuthKeySessionLayerStore) DeleteSessionLayer(
|
||||
ctx context.Context,
|
||||
rawAuthKeyID [8]byte,
|
||||
sessionID int64,
|
||||
) (bool, error) {
|
||||
return s.base.DeleteSessionLayer(ctx, rawAuthKeyID, sessionID)
|
||||
}
|
||||
|
||||
func (s *BatchedAuthKeySessionLayerStore) DeleteExpiredSessionLayers(ctx context.Context, limit int) (int, error) {
|
||||
return s.base.DeleteExpiredSessionLayers(ctx, limit)
|
||||
}
|
||||
|
||||
func (s *BatchedAuthKeySessionLayerStore) Close() {
|
||||
s.once.Do(func() {
|
||||
s.gate.Lock()
|
||||
s.closed = true
|
||||
close(s.stop)
|
||||
s.cancel()
|
||||
s.gate.Unlock()
|
||||
<-s.done
|
||||
})
|
||||
}
|
||||
|
||||
func (s *BatchedAuthKeySessionLayerStore) run(ctx context.Context) {
|
||||
defer close(s.done)
|
||||
pending := make([]authKeySessionLayerBatchRequest, 0, s.cfg.MaxSize)
|
||||
for {
|
||||
if len(pending) == 0 {
|
||||
select {
|
||||
case request := <-s.queue:
|
||||
pending = append(pending, request)
|
||||
case <-s.stop:
|
||||
s.failQueued(context.Canceled, pending)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
if len(pending) < s.cfg.MaxSize {
|
||||
timer := time.NewTimer(s.cfg.MaxWait)
|
||||
collect:
|
||||
for len(pending) < s.cfg.MaxSize {
|
||||
select {
|
||||
case request := <-s.queue:
|
||||
pending = append(pending, request)
|
||||
case <-timer.C:
|
||||
break collect
|
||||
case <-s.stop:
|
||||
if !timer.Stop() {
|
||||
select {
|
||||
case <-timer.C:
|
||||
default:
|
||||
}
|
||||
}
|
||||
s.failQueued(context.Canceled, pending)
|
||||
return
|
||||
}
|
||||
}
|
||||
if !timer.Stop() {
|
||||
select {
|
||||
case <-timer.C:
|
||||
default:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
batch, remaining := selectDistinctLayerAdvanceBatch(pending, s.cfg.MaxSize)
|
||||
pending = remaining
|
||||
s.execute(ctx, batch)
|
||||
}
|
||||
}
|
||||
|
||||
func selectDistinctLayerAdvanceBatch(
|
||||
pending []authKeySessionLayerBatchRequest,
|
||||
maxSize int,
|
||||
) ([]authKeySessionLayerBatchRequest, []authKeySessionLayerBatchRequest) {
|
||||
batch := make([]authKeySessionLayerBatchRequest, 0, min(maxSize, len(pending)))
|
||||
remaining := make([]authKeySessionLayerBatchRequest, 0, len(pending))
|
||||
seen := make(map[authKeySessionLayerBatchKey]struct{}, min(maxSize, len(pending)))
|
||||
for _, request := range pending {
|
||||
if len(batch) >= maxSize {
|
||||
remaining = append(remaining, request)
|
||||
continue
|
||||
}
|
||||
key := authKeySessionLayerBatchKey{rawAuthKeyID: request.rawAuthKeyID, sessionID: request.sessionID}
|
||||
if _, exists := seen[key]; exists {
|
||||
remaining = append(remaining, request)
|
||||
continue
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
batch = append(batch, request)
|
||||
}
|
||||
return batch, remaining
|
||||
}
|
||||
|
||||
func (s *BatchedAuthKeySessionLayerStore) execute(ctx context.Context, batch []authKeySessionLayerBatchRequest) {
|
||||
active := batch[:0]
|
||||
for _, request := range batch {
|
||||
if err := request.ctx.Err(); err != nil {
|
||||
request.result <- authKeySessionLayerBatchResult{err: err}
|
||||
continue
|
||||
}
|
||||
active = append(active, request)
|
||||
}
|
||||
if len(active) == 0 {
|
||||
return
|
||||
}
|
||||
queryCtx, cancel := context.WithTimeout(ctx, s.cfg.QueryTimeout)
|
||||
results, err := s.base.tryAdvanceSessionLayersSameLayer(queryCtx, active)
|
||||
cancel()
|
||||
if err != nil {
|
||||
for _, request := range active {
|
||||
request.result <- authKeySessionLayerBatchResult{err: err}
|
||||
}
|
||||
return
|
||||
}
|
||||
for index, request := range active {
|
||||
request.result <- results[index]
|
||||
}
|
||||
}
|
||||
|
||||
func (s *BatchedAuthKeySessionLayerStore) failQueued(err error, pending []authKeySessionLayerBatchRequest) {
|
||||
for _, request := range pending {
|
||||
request.result <- authKeySessionLayerBatchResult{err: err}
|
||||
}
|
||||
for {
|
||||
select {
|
||||
case request := <-s.queue:
|
||||
request.result <- authKeySessionLayerBatchResult{err: err}
|
||||
default:
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *AuthKeyStore) tryAdvanceSessionLayersSameLayer(
|
||||
ctx context.Context,
|
||||
requests []authKeySessionLayerBatchRequest,
|
||||
) ([]authKeySessionLayerBatchResult, error) {
|
||||
results := make([]authKeySessionLayerBatchResult, len(requests))
|
||||
if len(requests) == 0 {
|
||||
return results, nil
|
||||
}
|
||||
rawIDs := make([]int64, len(requests))
|
||||
sessionIDs := make([]int64, len(requests))
|
||||
layers := make([]int32, len(requests))
|
||||
msgIDs := make([]int64, len(requests))
|
||||
expiresAts := make([]time.Time, len(requests))
|
||||
seen := make(map[authKeySessionLayerBatchKey]struct{}, len(requests))
|
||||
for index, request := range requests {
|
||||
key := authKeySessionLayerBatchKey{rawAuthKeyID: request.rawAuthKeyID, sessionID: request.sessionID}
|
||||
if _, duplicate := seen[key]; duplicate {
|
||||
return nil, fmt.Errorf("advance same-Layer auth key session batch: duplicate identity at index %d", index)
|
||||
}
|
||||
seen[key] = struct{}{}
|
||||
rawIDs[index] = authKeyIDToInt64(request.rawAuthKeyID)
|
||||
sessionIDs[index] = request.sessionID
|
||||
layers[index] = int32(request.layer)
|
||||
msgIDs[index] = request.msgID
|
||||
expiresAts[index] = request.expiresAt
|
||||
}
|
||||
rows, err := s.db.Query(ctx, `
|
||||
WITH input AS (
|
||||
SELECT *
|
||||
FROM unnest(
|
||||
$1::bigint[],
|
||||
$2::bigint[],
|
||||
$3::integer[],
|
||||
$4::bigint[],
|
||||
$5::timestamptz[]
|
||||
) WITH ORDINALITY AS value(raw_id, session_id, layer, msg_id, expires_at, ordinal)
|
||||
), identity AS MATERIALIZED (
|
||||
SELECT input.*,
|
||||
defaults.layer AS default_layer,
|
||||
defaults.layer_observation_id AS default_observation_id
|
||||
FROM input
|
||||
JOIN auth_keys AS raw
|
||||
ON raw.auth_key_id = input.raw_id
|
||||
LEFT JOIN temp_auth_key_bindings AS binding
|
||||
ON binding.temp_auth_key_id = raw.auth_key_id
|
||||
JOIN auth_keys AS defaults
|
||||
ON defaults.auth_key_id = COALESCE(binding.perm_auth_key_id, raw.auth_key_id)
|
||||
WHERE binding.temp_auth_key_id IS NULL
|
||||
OR (raw.expires_at > 0 AND defaults.expires_at = 0)
|
||||
), candidates AS MATERIALIZED (
|
||||
SELECT identity.ordinal,
|
||||
identity.msg_id,
|
||||
identity.expires_at,
|
||||
identity.default_layer,
|
||||
identity.default_observation_id,
|
||||
evidence.raw_auth_key_id,
|
||||
evidence.session_id,
|
||||
evidence.layer,
|
||||
evidence.observation_id
|
||||
FROM identity
|
||||
JOIN auth_key_session_layers AS evidence
|
||||
ON evidence.raw_auth_key_id = identity.raw_id
|
||||
AND evidence.session_id = identity.session_id
|
||||
WHERE evidence.layer = identity.layer
|
||||
AND evidence.msg_id < identity.msg_id
|
||||
AND evidence.expires_at > now()
|
||||
AND identity.layer > 0
|
||||
AND identity.msg_id > 0
|
||||
AND identity.msg_id % 4 = 0
|
||||
AND (identity.msg_id & 4294967295) <> 0
|
||||
AND identity.expires_at > now()
|
||||
AND identity.expires_at - interval '301 seconds' <= now() + interval '30 seconds'
|
||||
ORDER BY evidence.raw_auth_key_id, evidence.session_id
|
||||
FOR UPDATE OF evidence
|
||||
), advanced AS (
|
||||
UPDATE auth_key_session_layers AS evidence
|
||||
SET msg_id = candidates.msg_id,
|
||||
expires_at = candidates.expires_at
|
||||
FROM candidates
|
||||
WHERE evidence.raw_auth_key_id = candidates.raw_auth_key_id
|
||||
AND evidence.session_id = candidates.session_id
|
||||
RETURNING candidates.ordinal,
|
||||
candidates.default_layer,
|
||||
candidates.default_observation_id,
|
||||
evidence.layer,
|
||||
evidence.msg_id,
|
||||
evidence.observation_id,
|
||||
evidence.expires_at
|
||||
)
|
||||
SELECT ordinal,
|
||||
layer,
|
||||
msg_id,
|
||||
observation_id,
|
||||
expires_at,
|
||||
default_layer = layer AND default_observation_id = observation_id
|
||||
FROM advanced
|
||||
ORDER BY ordinal
|
||||
`, rawIDs, sessionIDs, layers, msgIDs, expiresAts)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("advance same-Layer auth key session batch: %w", err)
|
||||
}
|
||||
defer rows.Close()
|
||||
for rows.Next() {
|
||||
var (
|
||||
ordinal int64
|
||||
current store.AuthKeySessionLayer
|
||||
)
|
||||
if err := rows.Scan(
|
||||
&ordinal,
|
||||
¤t.Layer,
|
||||
¤t.MessageID,
|
||||
¤t.ObservationID,
|
||||
¤t.ExpiresAt,
|
||||
¤t.SharedDefault,
|
||||
); err != nil {
|
||||
return nil, fmt.Errorf("scan same-Layer auth key session batch: %w", err)
|
||||
}
|
||||
index := int(ordinal - 1)
|
||||
if index < 0 || index >= len(results) || results[index].fast {
|
||||
return nil, fmt.Errorf("advance same-Layer auth key session batch: invalid ordinal %d", ordinal)
|
||||
}
|
||||
results[index] = authKeySessionLayerBatchResult{current: current, fast: true}
|
||||
}
|
||||
if err := rows.Err(); err != nil {
|
||||
return nil, fmt.Errorf("advance same-Layer auth key session batch rows: %w", err)
|
||||
}
|
||||
return results, nil
|
||||
}
|
||||
|
||||
var _ store.AuthKeySessionLayerStore = (*BatchedAuthKeySessionLayerStore)(nil)
|
||||
Loading…
Add table
Add a link
Reference in a new issue