Merge pull request #17 from bundleman/fix-cdp

fix(drivers): fix routines leakage
bundleman · Jan 27, 2022 · 1ebe3c9 · 1ebe3c9
2 parents 087116c + b7375ba
commit 1ebe3c9
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/pkg/drivers/cdp/network/streams.go b/pkg/drivers/cdp/network/streams.go
@@ -45,6 +45,7 @@ func newNavigationEventStream(
 func (s *NavigationEventStream) Read(ctx context.Context) <-chan rtEvents.Message {
 	ch := make(chan rtEvents.Message)
 
+	ctx, cancel := context.WithCancel(context.Background())
 	go func() {
 		defer close(ch)
 
@@ -54,16 +55,15 @@ func (s *NavigationEventStream) Read(ctx context.Context) <-chan rtEvents.Messag
 				return
 			case <-s.onDoc.Ready():
 				if ctx.Err() != nil {
-					return
+					cancel()
 				}
 
 				repl, err := s.onDoc.Recv()
-
 				if err != nil {
 					ch <- rtEvents.WithErr(err)
 					s.logger.Trace().Err(err).Msg("failed to read data from within document navigation event stream")
 
-					return
+					cancel()
 				}
 
 				evt := NavigationEvent{
@@ -80,9 +80,11 @@ func (s *NavigationEventStream) Read(ctx context.Context) <-chan rtEvents.Messag
 				s.tail.Store(evt)
 
 				ch <- rtEvents.WithValue(&evt)
+				cancel()
+
 			case <-s.onFrame.Ready():
 				if ctx.Err() != nil {
-					return
+					cancel()
 				}
 
 				repl, err := s.onFrame.Recv()
@@ -91,7 +93,7 @@ func (s *NavigationEventStream) Read(ctx context.Context) <-chan rtEvents.Messag
 					ch <- rtEvents.WithErr(err)
 					s.logger.Trace().Err(err).Msg("failed to read data from frame navigation event stream")
 
-					return
+					cancel()
 				}
 
 				evt := NavigationEvent{
@@ -109,6 +111,7 @@ func (s *NavigationEventStream) Read(ctx context.Context) <-chan rtEvents.Messag
 				s.tail.Store(evt)
 
 				ch <- rtEvents.WithValue(&evt)
+				cancel()
 			}
 		}
 	}()