diff --git a/persistence/inmemory/inmemory.go b/persistence/inmemory/inmemory.go index 31e4bb8a..45e752a1 100644 --- a/persistence/inmemory/inmemory.go +++ b/persistence/inmemory/inmemory.go @@ -31,16 +31,14 @@ import ( func New() *Persistence { return &Persistence{ checkpoints: make(map[string][]byte), - logs: make(map[string]omniwitness.Log), } } type Persistence struct { - // mu allows checkpoints to be read concurrently, but - // exclusively locked for writing. - mu sync.RWMutex - checkpoints map[string][]byte - logs map[string]omniwitness.Log + checkpointsMu sync.RWMutex + checkpoints map[string][]byte + + logs sync.Map } func (p *Persistence) Init(_ context.Context) error { @@ -48,48 +46,40 @@ func (p *Persistence) Init(_ context.Context) error { } func (p *Persistence) AddLogs(ctx context.Context, lc []omniwitness.Log) error { - p.mu.Lock() - defer p.mu.Unlock() for _, l := range lc { logID := log.ID(l.Origin) - p.logs[logID] = l + p.logs.Store(logID, l) } return nil } func (p *Persistence) Logs(ctx context.Context) iter.Seq2[omniwitness.Log, error] { - p.mu.RLock() - defer p.mu.RUnlock() return func(yield func(omniwitness.Log, error) bool) { - for _, lc := range p.logs { - if !yield(lc, nil) { - return - } - } + p.logs.Range(func(key, value any) bool { + return yield(value.(omniwitness.Log), nil) + }) } } func (p *Persistence) Log(ctx context.Context, origin string) (omniwitness.Log, bool, error) { - p.mu.RLock() - defer p.mu.RUnlock() logID := log.ID(origin) - lc, ok := p.logs[logID] + val, ok := p.logs.Load(logID) if !ok { - return lc, false, nil + return omniwitness.Log{}, false, nil } - return lc, true, nil + return val.(omniwitness.Log), true, nil } func (p *Persistence) Latest(_ context.Context, origin string) ([]byte, error) { - p.mu.RLock() - defer p.mu.RUnlock() + p.checkpointsMu.RLock() + defer p.checkpointsMu.RUnlock() logID := log.ID(origin) return p.checkpoints[logID], nil } func (p *Persistence) Update(_ context.Context, origin string, f func([]byte) ([]byte, error)) error { - p.mu.Lock() - defer p.mu.Unlock() + p.checkpointsMu.Lock() + defer p.checkpointsMu.Unlock() logID := log.ID(origin) u, err := f(p.checkpoints[logID]) if err != nil { diff --git a/persistence/inmemory/inmemory_test.go b/persistence/inmemory/inmemory_test.go index 91e407aa..190f64dc 100644 --- a/persistence/inmemory/inmemory_test.go +++ b/persistence/inmemory/inmemory_test.go @@ -58,3 +58,20 @@ func TestUpdateConcurrent(t *testing.T) { t.Errorf("expected at least one success but got %s", string(cp)) } } + +func TestUpdateDeadlock(t *testing.T) { + p := New() + origin := "foo" + err := p.Update(t.Context(), origin, func(current []byte) (next []byte, err error) { + // This calls Log which tries to acquire RLock while Update holds Lock. + _, _, err = p.Log(t.Context(), origin) + if err != nil { + return nil, err + } + return []byte(fmt.Sprintf("%s\nsuccess", origin)), nil + }) + if err != nil { + t.Fatal(err) + } +} +