74 lines
2.0 KiB
Go
74 lines
2.0 KiB
Go
package pipeline
|
|
|
|
import (
|
|
"context"
|
|
"fmt"
|
|
"sync"
|
|
"sync/atomic"
|
|
)
|
|
|
|
type memoryIdempotencyEntry struct {
|
|
state IdempotencyState
|
|
token string
|
|
}
|
|
|
|
// MemoryIdempotencyStore only protects concurrent deliveries in one process.
|
|
// Complete releases the lease so a later version with the same stable checkId can update.
|
|
type MemoryIdempotencyStore struct {
|
|
mu sync.Mutex
|
|
entries map[string]memoryIdempotencyEntry
|
|
next atomic.Uint64
|
|
}
|
|
|
|
func NewMemoryIdempotencyStore() *MemoryIdempotencyStore {
|
|
return &MemoryIdempotencyStore{entries: make(map[string]memoryIdempotencyEntry)}
|
|
}
|
|
|
|
func (s *MemoryIdempotencyStore) Acquire(ctx context.Context, key string) (IdempotencyLease, error) {
|
|
if err := ctx.Err(); err != nil {
|
|
return IdempotencyLease{}, err
|
|
}
|
|
if key == "" {
|
|
return IdempotencyLease{}, fmt.Errorf("idempotency key is required")
|
|
}
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
if entry, ok := s.entries[key]; ok {
|
|
if entry.state == IdempotencyAcquired {
|
|
return IdempotencyLease{State: IdempotencyInProgress}, nil
|
|
}
|
|
return IdempotencyLease{State: entry.state}, nil
|
|
}
|
|
token := fmt.Sprintf("lease-%d", s.next.Add(1))
|
|
s.entries[key] = memoryIdempotencyEntry{state: IdempotencyAcquired, token: token}
|
|
return IdempotencyLease{State: IdempotencyAcquired, Token: token}, nil
|
|
}
|
|
|
|
func (s *MemoryIdempotencyStore) Complete(ctx context.Context, key, token string) error {
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
entry, ok := s.entries[key]
|
|
if !ok || entry.state != IdempotencyAcquired || entry.token != token {
|
|
return fmt.Errorf("idempotency lease is missing or stale")
|
|
}
|
|
delete(s.entries, key)
|
|
return nil
|
|
}
|
|
|
|
func (s *MemoryIdempotencyStore) Release(ctx context.Context, key, token string) error {
|
|
if err := ctx.Err(); err != nil {
|
|
return err
|
|
}
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
entry, ok := s.entries[key]
|
|
if !ok || entry.state != IdempotencyAcquired || entry.token != token {
|
|
return fmt.Errorf("idempotency lease is missing or stale")
|
|
}
|
|
delete(s.entries, key)
|
|
return nil
|
|
}
|