fix: support When custom context

pancsta · Jan 15, 2024 · 2960dc1 · 2960dc1
1 parent 8b81dee
commit 2960dc1
Showing 1 changed file with 28 additions and 0 deletions.
diff --git a/pkg/machine/machine.go b/pkg/machine/machine.go
@@ -265,6 +265,20 @@ func (m *Machine) When(states []string, ctx context.Context) chan struct{} {
 			m.indexWhen[s] = append(m.indexWhen[s], binding)
 		}
 	}
+	go func() {
+		// dispose the binding on ctx.Done() and m.Ctx.Done()
+		select {
+		case <-ctx.Done():
+		case <-m.Ctx.Done():
+			m.activeStatesLock.Lock()
+			for _, s := range states {
+				if _, ok := m.indexWhen[s]; ok {
+					m.indexWhen[s] = lo.Without(m.indexWhen[s], binding)
+				}
+			}
+			m.activeStatesLock.Unlock()
+		}
+	}()
 	m.activeStatesLock.Unlock()
 
 	return ch
@@ -310,6 +324,20 @@ func (m *Machine) WhenNot(states []string, ctx context.Context) chan struct{} {
 			m.indexWhen[s] = append(m.indexWhen[s], binding)
 		}
 	}
+	go func() {
+		// dispose the binding on ctx.Done() and m.Ctx.Done()
+		select {
+		case <-ctx.Done():
+		case <-m.Ctx.Done():
+			m.activeStatesLock.Lock()
+			for _, s := range states {
+				if _, ok := m.indexWhen[s]; ok {
+					m.indexWhen[s] = lo.Without(m.indexWhen[s], binding)
+				}
+			}
+			m.activeStatesLock.Unlock()
+		}
+	}()
 	m.activeStatesLock.Unlock()
 
 	return ch