From 79346309c27232e44c460ff3c31c1cc765be19f0 Mon Sep 17 00:00:00 2001
From: sean yu <55464069+hexbabe@users.noreply.github.com>
Date: Fri, 30 Aug 2024 12:16:02 -0400
Subject: [PATCH] Patch gostream release bug for modules (#4330)

---
 gostream/media.go      |  4 ++++
 gostream/media_test.go | 22 ++++++++++++++++++++++
 2 files changed, 26 insertions(+)

diff --git a/gostream/media.go b/gostream/media.go
index d5c33cc4a25..d5e34d7dc4e 100644
--- a/gostream/media.go
+++ b/gostream/media.go
@@ -360,6 +360,10 @@ func (pc *producerConsumer[T, U]) stop() {
 	pc.consumerCond.L.Unlock()
 	pc.activeBackgroundWorkers.Wait()
 
+	pc.currentMu.Lock()
+	defer pc.currentMu.Unlock()
+	pc.current = nil
+
 	// reset
 	cancelCtx, cancel := context.WithCancel(WithMIMETypeHint(pc.rootCancelCtx, pc.mimeType))
 	pc.cancelCtxMu.Lock()
diff --git a/gostream/media_test.go b/gostream/media_test.go
index fe781543285..7b6c45fb5e9 100644
--- a/gostream/media_test.go
+++ b/gostream/media_test.go
@@ -191,3 +191,25 @@ func TestStreamMultipleConsumers(t *testing.T) {
 		test.That(t, wrappedImg.released.Load(), test.ShouldBeTrue)
 	}
 }
+
+func TestStreamWithoutNext(t *testing.T) {
+	colors := []*WrappedImage{createWrappedImage(t, rimage.Red)}
+
+	imgSource := &imageSource{WrappedImages: colors}
+	videoSrc := NewVideoSource(imgSource, prop.Video{})
+
+	// Start stream
+	stream, err := videoSrc.Stream(context.Background())
+	test.That(t, err, test.ShouldBeNil)
+	// Get one frame
+	_, release, err := stream.Next(context.Background())
+	test.That(t, err, test.ShouldBeNil)
+	release()
+	// Close stream
+	stream.Close(context.Background())
+
+	// Spin up stream and close without calling Next
+	stream, err = videoSrc.Stream(context.Background())
+	test.That(t, err, test.ShouldBeNil)
+	stream.Close(context.Background())
+}