update(Timeout): wrap async callback with Context.PreAsyncCall

b97tsk · Jul 4, 2024 · e18982c · e18982c
1 parent 83a06a3
commit e18982c
Showing 1 changed file with 2 additions and 10 deletions.
diff --git a/timeout.go b/timeout.go
@@ -58,15 +58,7 @@ func (ob timeoutObservable[T]) Subscribe(parent Context, o Observer[T]) {
 
 	x.Context.Store(c.Context)
 
-	if c.WaitGroup != nil {
-		c.WaitGroup.Add(1)
-	}
-
-	tm := time.AfterFunc(ob.First, func() {
-		if c.WaitGroup != nil {
-			defer c.WaitGroup.Done()
-		}
-
+	tm := time.AfterFunc(ob.First, c.PreAsyncCall(func() {
 		if x.Context.Swap(sentinel) != sentinel {
 			cancel()
 
@@ -77,7 +69,7 @@ func (ob timeoutObservable[T]) Subscribe(parent Context, o Observer[T]) {
 
 			o.Error(ErrTimeout)
 		}
-	})
+	}))
 
 	ob.Source.Subscribe(c, func(n Notification[T]) {
 		switch n.Kind {