refac(machine): extract When* ctx disposal

pancsta · Jul 10, 2024 · 7cc2191 · 7cc2191
1 parent ab1ee58
commit 7cc2191
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 86 deletions.
diff --git 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
@@ -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)
+				}
+			}
+		}
+	}()
+}