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/message_send_batch_actor.go
Normal file
408
internal/store/postgres/message_send_batch_actor.go
Normal file
|
|
@ -0,0 +1,408 @@
|
|||
package postgres
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/jackc/pgx/v5"
|
||||
"github.com/jackc/pgx/v5/pgxpool"
|
||||
|
||||
"telesrv/internal/domain"
|
||||
"telesrv/internal/store"
|
||||
"telesrv/internal/store/postgres/sqlcgen"
|
||||
)
|
||||
|
||||
const (
|
||||
plainPrivateSendBatchWorkers = 8
|
||||
plainPrivateSendBatchMaxTasks = 32
|
||||
plainPrivateSendBatchMinTasks = 8
|
||||
plainPrivateSendBatchMailbox = 8192
|
||||
plainPrivateSendBatchFlushInterval = 8 * time.Millisecond
|
||||
plainPrivateSendBatchTimeout = 10 * time.Second
|
||||
plainPrivateSendBatchMaxQueuedBytes = int64(64 << 20)
|
||||
plainPrivateSendTaskFixedBytes = int64(1024)
|
||||
)
|
||||
|
||||
var errPlainPrivateSendBatchOverloaded = errors.New("postgres: plain private send batch actor overloaded")
|
||||
|
||||
var processPlainPrivateSendBatcher = newPlainPrivateSendBatchActor()
|
||||
|
||||
type plainPrivateSendBatchTask struct {
|
||||
ctx context.Context
|
||||
store *MessageStore
|
||||
req domain.SendPrivateTextRequest
|
||||
fingerprint []byte
|
||||
bytes int64
|
||||
done chan plainPrivateSendBatchResult
|
||||
}
|
||||
|
||||
type plainPrivateSendBatchResult struct {
|
||||
result domain.SendPrivateTextResult
|
||||
err error
|
||||
}
|
||||
|
||||
type plainPrivateSendBatchCompletion struct {
|
||||
tasks []*plainPrivateSendBatchTask
|
||||
results []plainPrivateSendBatchResult
|
||||
}
|
||||
|
||||
type plainPrivateSendScope struct {
|
||||
store *MessageStore
|
||||
userID int64
|
||||
}
|
||||
|
||||
type plainPrivateSendBatchActor struct {
|
||||
submit chan *plainPrivateSendBatchTask
|
||||
work chan []*plainPrivateSendBatchTask
|
||||
completed chan plainPrivateSendBatchCompletion
|
||||
queuedBytes atomic.Int64
|
||||
batches atomic.Uint64
|
||||
tasks atomic.Uint64
|
||||
}
|
||||
|
||||
type plainPrivateSendBatchSnapshot struct {
|
||||
Batches uint64
|
||||
Tasks uint64
|
||||
}
|
||||
|
||||
func newPlainPrivateSendBatchActor() *plainPrivateSendBatchActor {
|
||||
a := &plainPrivateSendBatchActor{
|
||||
submit: make(chan *plainPrivateSendBatchTask, plainPrivateSendBatchMailbox),
|
||||
work: make(chan []*plainPrivateSendBatchTask, plainPrivateSendBatchWorkers),
|
||||
completed: make(chan plainPrivateSendBatchCompletion, plainPrivateSendBatchWorkers),
|
||||
}
|
||||
for range plainPrivateSendBatchWorkers {
|
||||
go a.worker()
|
||||
}
|
||||
go a.run()
|
||||
return a
|
||||
}
|
||||
|
||||
func (a *plainPrivateSendBatchActor) Snapshot() plainPrivateSendBatchSnapshot {
|
||||
if a == nil {
|
||||
return plainPrivateSendBatchSnapshot{}
|
||||
}
|
||||
return plainPrivateSendBatchSnapshot{Batches: a.batches.Load(), Tasks: a.tasks.Load()}
|
||||
}
|
||||
|
||||
func (a *plainPrivateSendBatchActor) Eligible(messageStore *MessageStore) bool {
|
||||
if a == nil || messageStore == nil || messageStore.boxIDs == nil {
|
||||
return false
|
||||
}
|
||||
if _, ok := messageStore.db.(*pgxpool.Pool); !ok {
|
||||
return false
|
||||
}
|
||||
// The process batch path is a distributed allocator + PostgreSQL boundary.
|
||||
// Local and test allocators cannot reserve gap-safe ids before the
|
||||
// transaction under cross-process concurrency and therefore remain on the
|
||||
// single-command transaction path.
|
||||
_, distributed := messageStore.boxIDs.(store.DistributedBoxIDAllocator)
|
||||
return distributed
|
||||
}
|
||||
|
||||
func (a *plainPrivateSendBatchActor) Submit(
|
||||
ctx context.Context,
|
||||
messageStore *MessageStore,
|
||||
req domain.SendPrivateTextRequest,
|
||||
fingerprint []byte,
|
||||
) (domain.SendPrivateTextResult, error) {
|
||||
if !a.Eligible(messageStore) || ctx == nil || len(fingerprint) == 0 {
|
||||
return domain.SendPrivateTextResult{}, fmt.Errorf("postgres: invalid plain private send batch submission")
|
||||
}
|
||||
retained := plainPrivateSendTaskFixedBytes + int64(len(req.Message)+len(fingerprint))
|
||||
if !reservePlainPrivateSendBatchBytes(&a.queuedBytes, retained) {
|
||||
return domain.SendPrivateTextResult{}, errPlainPrivateSendBatchOverloaded
|
||||
}
|
||||
task := &plainPrivateSendBatchTask{
|
||||
ctx: ctx, store: messageStore, req: req,
|
||||
fingerprint: append([]byte(nil), fingerprint...),
|
||||
bytes: retained,
|
||||
done: make(chan plainPrivateSendBatchResult, 1),
|
||||
}
|
||||
select {
|
||||
case a.submit <- task:
|
||||
case <-ctx.Done():
|
||||
a.queuedBytes.Add(-retained)
|
||||
return domain.SendPrivateTextResult{}, ctx.Err()
|
||||
}
|
||||
select {
|
||||
case result := <-task.done:
|
||||
return result.result, result.err
|
||||
case <-ctx.Done():
|
||||
// Once accepted, the bounded actor may commit after the caller stops
|
||||
// waiting. random_id is the durable receipt for an exact replay.
|
||||
return domain.SendPrivateTextResult{}, ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
func reservePlainPrivateSendBatchBytes(used *atomic.Int64, amount int64) bool {
|
||||
if used == nil || amount <= 0 || amount > plainPrivateSendBatchMaxQueuedBytes {
|
||||
return false
|
||||
}
|
||||
for {
|
||||
current := used.Load()
|
||||
if current > plainPrivateSendBatchMaxQueuedBytes-amount {
|
||||
return false
|
||||
}
|
||||
if used.CompareAndSwap(current, current+amount) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (a *plainPrivateSendBatchActor) run() {
|
||||
ticker := time.NewTicker(plainPrivateSendBatchFlushInterval)
|
||||
defer ticker.Stop()
|
||||
pending := make([]*plainPrivateSendBatchTask, 0, plainPrivateSendBatchMailbox)
|
||||
busy := make(map[plainPrivateSendScope]struct{})
|
||||
available := plainPrivateSendBatchWorkers
|
||||
|
||||
completeCanceled := func(task *plainPrivateSendBatchTask) {
|
||||
err := context.Canceled
|
||||
if task.ctx != nil && task.ctx.Err() != nil {
|
||||
err = task.ctx.Err()
|
||||
}
|
||||
task.done <- plainPrivateSendBatchResult{err: err}
|
||||
a.queuedBytes.Add(-task.bytes)
|
||||
}
|
||||
dispatch := func(flush bool) {
|
||||
for available > 0 && len(pending) > 0 {
|
||||
if !flush && len(pending) < plainPrivateSendBatchMinTasks {
|
||||
return
|
||||
}
|
||||
batch, remaining, canceled := selectPlainPrivateSendBatch(pending, busy)
|
||||
pending = remaining
|
||||
for _, task := range canceled {
|
||||
completeCanceled(task)
|
||||
}
|
||||
if len(batch) == 0 {
|
||||
return
|
||||
}
|
||||
for _, task := range batch {
|
||||
for _, userID := range plainPrivateSendTaskUsers(task) {
|
||||
busy[plainPrivateSendScope{store: task.store, userID: userID}] = struct{}{}
|
||||
}
|
||||
}
|
||||
available--
|
||||
a.work <- batch
|
||||
}
|
||||
}
|
||||
|
||||
for {
|
||||
select {
|
||||
case task := <-a.submit:
|
||||
pending = append(pending, task)
|
||||
dispatch(false)
|
||||
case completion := <-a.completed:
|
||||
available++
|
||||
for i, task := range completion.tasks {
|
||||
for _, userID := range plainPrivateSendTaskUsers(task) {
|
||||
delete(busy, plainPrivateSendScope{store: task.store, userID: userID})
|
||||
}
|
||||
result := plainPrivateSendBatchResult{err: errors.New("postgres: missing plain private send batch result")}
|
||||
if i < len(completion.results) {
|
||||
result = completion.results[i]
|
||||
}
|
||||
task.done <- result
|
||||
a.queuedBytes.Add(-task.bytes)
|
||||
}
|
||||
dispatch(false)
|
||||
case <-ticker.C:
|
||||
dispatch(true)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func selectPlainPrivateSendBatch(
|
||||
pending []*plainPrivateSendBatchTask,
|
||||
busy map[plainPrivateSendScope]struct{},
|
||||
) (batch, remaining, canceled []*plainPrivateSendBatchTask) {
|
||||
selected := make(map[plainPrivateSendScope]struct{}, plainPrivateSendBatchMaxTasks*2)
|
||||
blockedByOlder := make(map[plainPrivateSendScope]struct{}, plainPrivateSendBatchMaxTasks*2)
|
||||
var selectedStore *MessageStore
|
||||
remaining = make([]*plainPrivateSendBatchTask, 0, len(pending))
|
||||
for _, task := range pending {
|
||||
if task == nil || task.ctx == nil || task.ctx.Err() != nil {
|
||||
if task != nil {
|
||||
canceled = append(canceled, task)
|
||||
}
|
||||
continue
|
||||
}
|
||||
users := plainPrivateSendTaskUsers(task)
|
||||
blocked := selectedStore != nil && task.store != selectedStore
|
||||
for _, userID := range users {
|
||||
key := plainPrivateSendScope{store: task.store, userID: userID}
|
||||
if _, ok := busy[key]; ok {
|
||||
blocked = true
|
||||
}
|
||||
if _, ok := blockedByOlder[key]; ok {
|
||||
blocked = true
|
||||
}
|
||||
if _, ok := selected[key]; ok {
|
||||
blocked = true
|
||||
}
|
||||
}
|
||||
if blocked || len(batch) >= plainPrivateSendBatchMaxTasks {
|
||||
remaining = append(remaining, task)
|
||||
for _, userID := range users {
|
||||
blockedByOlder[plainPrivateSendScope{store: task.store, userID: userID}] = struct{}{}
|
||||
}
|
||||
continue
|
||||
}
|
||||
if selectedStore == nil {
|
||||
selectedStore = task.store
|
||||
}
|
||||
batch = append(batch, task)
|
||||
for _, userID := range users {
|
||||
selected[plainPrivateSendScope{store: task.store, userID: userID}] = struct{}{}
|
||||
}
|
||||
}
|
||||
return batch, remaining, canceled
|
||||
}
|
||||
|
||||
func plainPrivateSendTaskUsers(task *plainPrivateSendBatchTask) []int64 {
|
||||
if task == nil {
|
||||
return nil
|
||||
}
|
||||
return normalizedUserLaneIDs([]int64{task.req.SenderUserID, task.req.RecipientUserID})
|
||||
}
|
||||
|
||||
func (a *plainPrivateSendBatchActor) worker() {
|
||||
for batch := range a.work {
|
||||
results := executePlainPrivateSendBatch(batch)
|
||||
a.batches.Add(1)
|
||||
a.tasks.Add(uint64(len(batch)))
|
||||
a.completed <- plainPrivateSendBatchCompletion{tasks: batch, results: results}
|
||||
}
|
||||
}
|
||||
|
||||
func executePlainPrivateSendBatch(tasks []*plainPrivateSendBatchTask) []plainPrivateSendBatchResult {
|
||||
results := make([]plainPrivateSendBatchResult, len(tasks))
|
||||
if len(tasks) == 0 || tasks[0] == nil || tasks[0].store == nil {
|
||||
return plainPrivateSendBatchErrorResults(results, errors.New("postgres: empty plain private send batch"))
|
||||
}
|
||||
messageStore := tasks[0].store
|
||||
pool, ok := messageStore.db.(*pgxpool.Pool)
|
||||
if !ok {
|
||||
return plainPrivateSendBatchErrorResults(results, errors.New("postgres: plain private send batch requires pgx pool"))
|
||||
}
|
||||
allocationUsers := make([]int64, 0, len(tasks)*2)
|
||||
lockUsers := make([]int64, 0, len(tasks)*2)
|
||||
seen := make(map[int64]struct{}, len(tasks)*2)
|
||||
for _, task := range tasks {
|
||||
if task == nil || task.store != messageStore {
|
||||
return plainPrivateSendBatchErrorResults(results, errors.New("postgres: mixed message stores in plain private send batch"))
|
||||
}
|
||||
for _, userID := range plainPrivateSendTaskUsers(task) {
|
||||
if _, exists := seen[userID]; exists {
|
||||
return plainPrivateSendBatchErrorResults(results, fmt.Errorf("postgres: overlapping user %d in plain private send batch", userID))
|
||||
}
|
||||
seen[userID] = struct{}{}
|
||||
}
|
||||
lockUsers = append(lockUsers, task.req.SenderUserID, task.req.RecipientUserID)
|
||||
allocationUsers = append(allocationUsers, task.req.SenderUserID)
|
||||
if task.req.SenderUserID != task.req.RecipientUserID && !task.req.RecipientBlocked {
|
||||
allocationUsers = append(allocationUsers, task.req.RecipientUserID)
|
||||
}
|
||||
}
|
||||
lockUsers = normalizedUserLaneIDs(lockUsers)
|
||||
allocationUsers = normalizedUserLaneIDs(allocationUsers)
|
||||
parent := context.Background()
|
||||
if tasks[0].ctx != nil {
|
||||
parent = context.WithoutCancel(tasks[0].ctx)
|
||||
}
|
||||
ctx, cancel := context.WithTimeout(parent, plainPrivateSendBatchTimeout)
|
||||
defer cancel()
|
||||
|
||||
boxIDs, err := messageStore.boxIDs.NextBoxIDs(ctx, allocationUsers)
|
||||
if err != nil {
|
||||
return plainPrivateSendBatchErrorResults(results, fmt.Errorf("allocate plain private send batch box ids: %w", err))
|
||||
}
|
||||
releaseLanes, err := messageStore.privateSendLanes.acquire(ctx, lockUsers...)
|
||||
if err != nil {
|
||||
return plainPrivateSendBatchErrorResults(results, fmt.Errorf("admit plain private send batch: %w", err))
|
||||
}
|
||||
defer releaseLanes()
|
||||
|
||||
tx, err := pool.Begin(ctx)
|
||||
if err != nil {
|
||||
return plainPrivateSendBatchErrorResults(results, fmt.Errorf("begin plain private send batch: %w", err))
|
||||
}
|
||||
committed := false
|
||||
defer func() {
|
||||
if !committed {
|
||||
_ = tx.Rollback(context.Background())
|
||||
}
|
||||
}()
|
||||
if err := lockUsersForUpdate(ctx, tx, lockUsers...); err != nil {
|
||||
return plainPrivateSendBatchErrorResults(results, fmt.Errorf("lock plain private send batch users: %w", err))
|
||||
}
|
||||
if err := lockDispatchOutboxAppendFences(ctx, tx, allocationUsers); err != nil {
|
||||
return plainPrivateSendBatchErrorResults(results, err)
|
||||
}
|
||||
created, err := createPlainPrivateMessageBatch(ctx, tx, tasks)
|
||||
if err != nil {
|
||||
return plainPrivateSendBatchErrorResults(results, err)
|
||||
}
|
||||
qtx := sqlcgen.New(tx)
|
||||
for i, task := range tasks {
|
||||
if created[i].inserted {
|
||||
continue
|
||||
}
|
||||
duplicate, found, duplicateErr := messageStore.duplicateSendResult(ctx, qtx, task.req, task.fingerprint)
|
||||
if duplicateErr != nil {
|
||||
return plainPrivateSendBatchErrorResults(results, duplicateErr)
|
||||
}
|
||||
if !found {
|
||||
return plainPrivateSendBatchErrorResults(results, errors.New("duplicate batched private message disappeared after unique conflict"))
|
||||
}
|
||||
duplicate.Duplicate = true
|
||||
results[i].result = duplicate
|
||||
}
|
||||
projections, err := persistPlainPrivateSendProjectionBatch(ctx, tx, tasks, created, boxIDs)
|
||||
if err != nil {
|
||||
return plainPrivateSendBatchErrorResults(results, err)
|
||||
}
|
||||
for i := range tasks {
|
||||
if !created[i].inserted {
|
||||
continue
|
||||
}
|
||||
projection, ok := projections[i]
|
||||
if !ok {
|
||||
return plainPrivateSendBatchErrorResults(results, fmt.Errorf("missing plain private send batch projection %d", i))
|
||||
}
|
||||
results[i].result = domain.SendPrivateTextResult{
|
||||
SenderMessage: projection.Sender, RecipientMessage: projection.Recipient,
|
||||
SenderEvent: eventFromMessage(projection.Sender), RecipientEvent: eventFromMessage(projection.Recipient),
|
||||
}
|
||||
}
|
||||
if err := tx.Commit(ctx); err != nil {
|
||||
return plainPrivateSendBatchErrorResults(results, fmt.Errorf("commit plain private send batch: %w", err))
|
||||
}
|
||||
committed = true
|
||||
return results
|
||||
}
|
||||
|
||||
func lockDispatchOutboxAppendFences(ctx context.Context, tx pgx.Tx, userIDs []int64) error {
|
||||
userIDs = normalizedUserLaneIDs(userIDs)
|
||||
if len(userIDs) == 0 {
|
||||
return nil
|
||||
}
|
||||
if _, err := tx.Exec(ctx, `
|
||||
SELECT pg_advisory_xact_lock_shared(dispatch_outbox_lane_advisory_key(streams.target_user_id))
|
||||
FROM unnest($1::bigint[]) AS streams(target_user_id)
|
||||
ORDER BY streams.target_user_id`, userIDs); err != nil {
|
||||
return fmt.Errorf("lock dispatch outbox append fences: %w", err)
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func plainPrivateSendBatchErrorResults(results []plainPrivateSendBatchResult, err error) []plainPrivateSendBatchResult {
|
||||
for i := range results {
|
||||
results[i] = plainPrivateSendBatchResult{err: err}
|
||||
}
|
||||
return results
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue