[FIXED] Ensure object watcher stop closes the updates channel (#1844)

Signed-off-by: Piotr Piotrowski <piotr@synadia.com>

jetstream/object.go

@@ -1330,6 +1330,9 @@ func (obs *obs) Watch(ctx context.Context, opts ...WatchOpt) (ObjectWatcher, err
 	if err != nil {
 		return nil, err
 	}
+	sub.SetClosedHandler(func(_ string) {
+		close(w.updates)
+	})
 	w.sub = sub
 	return w, nil
 }

jetstream/test/kv_test.go

@@ -614,6 +614,37 @@ func TestKeyValueWatch(t *testing.T) {
 		expectOk(t, kv.Delete(ctx, "age"))
 		expectDelete("age", 6)
 	})
+
+	t.Run("stop watcher should not block", func(t *testing.T) {
+		s := RunBasicJetStreamServer()
+		defer shutdownJSServerAndRemoveStorage(t, s)
+
+		nc, js := jsClient(t, s)
+		defer nc.Close()
+		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+
+		kv, err := js.CreateKeyValue(ctx, jetstream.KeyValueConfig{Bucket: "WATCH"})
+		expectOk(t, err)
+
+		watcher, err := kv.WatchAll(ctx)
+		expectOk(t, err)
+
+		expectInitDone := expectInitDoneF(t, watcher)
+		expectInitDone()
+
+		err = watcher.Stop()
+		expectOk(t, err)
+
+		select {
+		case _, ok := <-watcher.Updates():
+			if ok {
+				t.Fatalf("Expected channel to be closed")
+			}
+		case <-time.After(100 * time.Millisecond):
+			break
+		}
+	})
 }
 
 func TestKeyValueWatchContext(t *testing.T) {

jetstream/test/object_test.go

@@ -695,6 +695,36 @@ func TestObjectWatch(t *testing.T) {
 		expectOk(t, err)
 		expectUpdate("C")
 	})
+
+	t.Run("stop watcher should close the channel", func(t *testing.T) {
+		s := RunBasicJetStreamServer()
+		defer shutdownJSServerAndRemoveStorage(t, s)
+
+		nc, js := jsClient(t, s)
+		defer nc.Close()
+		ctx := context.Background()
+
+		obs, err := js.CreateObjectStore(ctx, jetstream.ObjectStoreConfig{Bucket: "WATCH-TEST"})
+		expectOk(t, err)
+
+		watcher, err := obs.Watch(ctx)
+		expectOk(t, err)
+
+		expectInitDone := expectInitDoneF(t, watcher)
+		expectInitDone()
+
+		err = watcher.Stop()
+		expectOk(t, err)
+
+		select {
+		case _, ok := <-watcher.Updates():
+			if ok {
+				t.Fatalf("Expected channel to be closed")
+			}
+		case <-time.After(100 * time.Millisecond):
+			return
+		}
+	})
 }
 
 func TestObjectLinks(t *testing.T) {

object.go

@@ -1127,6 +1127,10 @@ func (obs *obs) Watch(opts ...WatchOpt) (ObjectWatcher, error) {
 	if err != nil {
 		return nil, err
 	}
+	// Set us up to close when the waitForMessages func returns.
+	sub.pDone = func(_ string) {
+		close(w.updates)
+	}
 	w.sub = sub
 	return w, nil
 }

test/kv_test.go

@@ -454,6 +454,35 @@ func TestKeyValueWatch(t *testing.T) {
 		expectOk(t, kv.Delete("age"))
 		expectDelete("age", 6)
 	})
+
+	t.Run("stop watcher should not block", func(t *testing.T) {
+		s := RunBasicJetStreamServer()
+		defer shutdownJSServerAndRemoveStorage(t, s)
+
+		nc, js := jsClient(t, s)
+		defer nc.Close()
+
+		kv, err := js.CreateKeyValue(&nats.KeyValueConfig{Bucket: "WATCH"})
+		expectOk(t, err)
+
+		watcher, err := kv.WatchAll()
+		expectOk(t, err)
+
+		expectInitDone := expectInitDoneF(t, watcher)
+		expectInitDone()
+
+		err = watcher.Stop()
+		expectOk(t, err)
+
+		select {
+		case _, ok := <-watcher.Updates():
+			if ok {
+				t.Fatalf("Expected channel to be closed")
+			}
+		case <-time.After(100 * time.Millisecond):
+			t.Fatalf("Stop watcher did not return")
+		}
+	})
 }
 
 func TestKeyValueWatchContext(t *testing.T) {

test/object_test.go

@@ -625,6 +625,35 @@ func TestObjectWatch(t *testing.T) {
 		expectOk(t, err)
 		expectUpdate("C")
 	})
+
+	t.Run("stop watcher should not block", func(t *testing.T) {
+		s := RunBasicJetStreamServer()
+		defer shutdownJSServerAndRemoveStorage(t, s)
+
+		nc, js := jsClient(t, s)
+		defer nc.Close()
+
+		obs, err := js.CreateObjectStore(&nats.ObjectStoreConfig{Bucket: "WATCH-TEST"})
+		expectOk(t, err)
+
+		watcher, err := obs.Watch()
+		expectOk(t, err)
+
+		expectInitDone := expectInitDoneF(t, watcher)
+		expectInitDone()
+
+		err = watcher.Stop()
+		expectOk(t, err)
+
+		select {
+		case _, ok := <-watcher.Updates():
+			if ok {
+				t.Fatal("Expected channel to be closed")
+			}
+		case <-time.After(100 * time.Millisecond):
+			return
+		}
+	})
 }
 
 func TestObjectLinks(t *testing.T) {