diff --git a/pkg/machine/machine.go b/pkg/machine/machine.go index f0b6226..33c6f5e 100644 --- a/pkg/machine/machine.go +++ b/pkg/machine/machine.go @@ -90,7 +90,7 @@ type Machine struct { indexWhen indexWhen indexWhenTime indexWhenTime indexWhenArgs indexWhenArgs - indexWhenArgsLock sync.Mutex + indexWhenArgsLock sync.RWMutex indexStateCtx indexStateCtx indexEventCh indexEventCh indexEventChLock sync.Mutex @@ -533,35 +533,7 @@ func (m *Machine) WhenNot(states S, ctx context.Context) chan struct{} { } // dispose with context - // TODO extract - if ctx != nil { - go func() { - select { - case <-ch: - return - case <-m.Ctx.Done(): - return - case <-ctx.Done(): - } - // GC only if needed - if m.Disposed { - return - } - - m.activeStatesLock.Lock() - defer m.activeStatesLock.Unlock() - - for _, s := range states { - if _, ok := m.indexWhen[s]; ok { - if len(m.indexWhen[s]) == 1 { - delete(m.indexWhen, s) - } else { - m.indexWhen[s] = slicesWithout(m.indexWhen[s], binding) - } - } - } - }() - } + diposeWithCtx(m, ctx, ch, states, binding, &m.activeStatesLock, m.indexWhen) // insert the binding for _, s := range states { @@ -620,33 +592,8 @@ func (m *Machine) WhenArgs( } // dispose with context - // TODO extract - if ctx != nil { - go func() { - select { - case <-ch: - return - case <-m.Ctx.Done(): - return - case <-ctx.Done(): - } - // GC only if needed - if m.Disposed { - return - } - - m.indexWhenArgsLock.Lock() - defer m.indexWhenArgsLock.Unlock() - - if _, ok := m.indexWhenArgs[name]; ok { - if len(m.indexWhenArgs[name]) == 1 { - delete(m.indexWhenArgs, name) - } else { - m.indexWhenArgs[name] = slicesWithout(m.indexWhenArgs[name], binding) - } - } - }() - } + diposeWithCtx(m, ctx, ch, S{name}, binding, &m.indexWhenArgsLock, + m.indexWhenArgs) // insert the binding if _, ok := m.indexWhen[name]; !ok { @@ -715,35 +662,8 @@ func (m *Machine) WhenTime( } // dispose with context - // TODO extract - if ctx != nil { - go func() { - select { - case <-ch: - return - case <-m.Ctx.Done(): - return - case <-ctx.Done(): - } - // GC only if needed - if m.Disposed { - return - } - - m.activeStatesLock.Lock() - defer m.activeStatesLock.Unlock() - - for _, s := range states { - if _, ok := indexWhenTime[s]; ok { - if len(m.indexWhenTime[s]) == 1 { - delete(m.indexWhenTime, s) - } else { - m.indexWhenTime[s] = slicesWithout(m.indexWhenTime[s], binding) - } - } - } - }() - } + diposeWithCtx(m, ctx, ch, states, binding, &m.activeStatesLock, + m.indexWhenTime) // insert the binding for _, s := range states { diff --git a/pkg/machine/misc.go b/pkg/machine/misc.go index 16be728..b9a604e 100644 --- a/pkg/machine/misc.go +++ b/pkg/machine/misc.go @@ -10,6 +10,7 @@ import ( "regexp" "slices" "strings" + "sync" "time" ) @@ -782,3 +783,43 @@ func slicesUniq[S ~[]E, E comparable](coll S) S { } return ret } + +// diposeWithCtx handles early binding disposal caused by a canceled context. +// It's used by most of "when" methods. +func diposeWithCtx[T comparable]( + mach *Machine, ctx context.Context, ch chan struct{}, states S, binding T, + lock *sync.RWMutex, index map[string][]T, +) { + if ctx == nil { + return + } + go func() { + select { + case <-ch: + return + case <-mach.Ctx.Done(): + return + case <-ctx.Done(): + } + // GC only if needed + if mach.Disposed { + return + } + + // TODO track + closeSafe(ch) + + lock.Lock() + defer lock.Unlock() + + for _, s := range states { + if _, ok := index[s]; ok { + if len(index[s]) == 1 { + delete(index, s) + } else { + index[s] = slicesWithout(index[s], binding) + } + } + } + }() +}