96 lines
2.3 KiB
Go
96 lines
2.3 KiB
Go
package queue
|
|
|
|
import (
|
|
"context"
|
|
"sync"
|
|
"time"
|
|
)
|
|
|
|
type IdempotencyState string
|
|
|
|
const (
|
|
IdempotencyStateAcquired IdempotencyState = "acquired"
|
|
IdempotencyStateDuplicate IdempotencyState = "duplicate"
|
|
IdempotencyStateInProgress IdempotencyState = "in_progress"
|
|
)
|
|
|
|
type IdempotencyStore interface {
|
|
Acquire(ctx context.Context, key string, processingTTL time.Duration) (IdempotencyState, error)
|
|
MarkCompleted(ctx context.Context, key string, ttl time.Duration) error
|
|
Release(ctx context.Context, key string) error
|
|
}
|
|
|
|
type NoopIdempotencyStore struct{}
|
|
|
|
func (NoopIdempotencyStore) Acquire(context.Context, string, time.Duration) (IdempotencyState, error) {
|
|
return IdempotencyStateAcquired, nil
|
|
}
|
|
|
|
func (NoopIdempotencyStore) MarkCompleted(context.Context, string, time.Duration) error {
|
|
return nil
|
|
}
|
|
|
|
func (NoopIdempotencyStore) Release(context.Context, string) error {
|
|
return nil
|
|
}
|
|
|
|
type MemoryIdempotencyStore struct {
|
|
mu sync.Mutex
|
|
records map[string]memoryIdempotencyRecord
|
|
now func() time.Time
|
|
}
|
|
|
|
type memoryIdempotencyRecord struct {
|
|
state IdempotencyState
|
|
expiresAt time.Time
|
|
}
|
|
|
|
func NewMemoryIdempotencyStore() *MemoryIdempotencyStore {
|
|
return &MemoryIdempotencyStore{
|
|
records: make(map[string]memoryIdempotencyRecord),
|
|
now: func() time.Time { return time.Now().UTC() },
|
|
}
|
|
}
|
|
|
|
func (s *MemoryIdempotencyStore) Acquire(_ context.Context, key string, processingTTL time.Duration) (IdempotencyState, error) {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
s.cleanupExpiredLocked()
|
|
if record, ok := s.records[key]; ok {
|
|
return record.state, nil
|
|
}
|
|
s.records[key] = memoryIdempotencyRecord{
|
|
state: IdempotencyStateInProgress,
|
|
expiresAt: s.now().Add(processingTTL),
|
|
}
|
|
return IdempotencyStateAcquired, nil
|
|
}
|
|
|
|
func (s *MemoryIdempotencyStore) MarkCompleted(_ context.Context, key string, ttl time.Duration) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
|
|
s.records[key] = memoryIdempotencyRecord{
|
|
state: IdempotencyStateDuplicate,
|
|
expiresAt: s.now().Add(ttl),
|
|
}
|
|
return nil
|
|
}
|
|
|
|
func (s *MemoryIdempotencyStore) Release(_ context.Context, key string) error {
|
|
s.mu.Lock()
|
|
defer s.mu.Unlock()
|
|
delete(s.records, key)
|
|
return nil
|
|
}
|
|
|
|
func (s *MemoryIdempotencyStore) cleanupExpiredLocked() {
|
|
now := s.now()
|
|
for key, record := range s.records {
|
|
if !record.expiresAt.IsZero() && now.After(record.expiresAt) {
|
|
delete(s.records, key)
|
|
}
|
|
}
|
|
}
|