Merge pull request #531 from nats-io/fix-consumer-drain

[FIX] connection drain for consume messages iterators

jetstream/consumer.ts

@@ -36,6 +36,7 @@ import {
   PullOptions,
   ReplayPolicy,
 } from "./jsapi_types.ts";
+import { SubscriptionImpl } from "../nats-base-client/protocol.ts";
 
 enum PullConsumerType {
   Unset = -1,
@@ -371,6 +372,16 @@ export class PullConsumerMessagesImpl extends QueuedIteratorImpl<JsMsg>
       },
     });
 
+    this.sub.closed.then(() => {
+      // for ordered consumer we cannot break the iterator
+      if ((this.sub as SubscriptionImpl).draining) {
+        // @ts-ignore: we are pushing the pull fn
+        this._push(() => {
+          this.stop();
+        });
+      }
+    });
+
     if (idle_heartbeat) {
       this.monitor = new core.IdleHeartbeat(idle_heartbeat, (data): boolean => {
         // for the pull consumer - missing heartbeats may be corrected

jetstream/tests/consumers_test.ts

@@ -48,6 +48,7 @@ import {
   ConsumerStatus,
   PullConsumerMessagesImpl,
 } from "../consumer.ts";
+import { deadline } from "../../nats-base-client/util.ts";
 
 Deno.test("consumers - min supported server", async () => {
   const { ns, nc } = await setup(jetstreamServerConf({}));
@@ -915,3 +916,30 @@ Deno.test("consumers - sub leaks consume()", async () => {
   assertEquals(nc.protocol.subscriptions.size(), 1);
   await cleanup(ns, nc);
 });
+
+Deno.test("consumers - consume drain", async () => {
+  const { ns, nc } = await setup(jetstreamServerConf());
+  const { stream } = await initStream(nc);
+
+  const jsm = await nc.jetstreamManager();
+  await jsm.consumers.add(stream, {
+    durable_name: stream,
+    ack_policy: AckPolicy.Explicit,
+  });
+  //@ts-ignore: test
+  const js = nc.jetstream();
+  const c = await js.consumers.get(stream, stream);
+  const iter = await c.consume({ expires: 30000 });
+  setTimeout(() => {
+    nc.drain();
+  }, 100);
+  const done = (async () => {
+    for await (const _m of iter) {
+      // nothing
+    }
+  })().then();
+
+  await deadline(done, 1000);
+
+  await cleanup(ns, nc);
+});

jetstream/tests/consumersordered_test.ts

@@ -28,7 +28,7 @@ import {
   notCompatible,
   setup,
 } from "../../tests/helpers/mod.ts";
-import { delay } from "../../nats-base-client/util.ts";
+import { deadline, delay } from "../../nats-base-client/util.ts";
 
 Deno.test("ordered - get", async () => {
   const { ns, nc } = await setup(jetstreamServerConf());
@@ -632,3 +632,24 @@ Deno.test("ordered - sub leaks consume()", async () => {
   assertEquals(nc.protocol.subscriptions.size(), 1);
   await cleanup(ns, nc);
 });
+
+Deno.test("ordered - consume drain", async () => {
+  const { ns, nc } = await setup(jetstreamServerConf());
+  const { stream } = await initStream(nc);
+
+  const js = nc.jetstream();
+  const c = await js.consumers.get(stream);
+  const iter = await c.consume({ expires: 30000 });
+  setTimeout(() => {
+    nc.drain();
+  }, 100);
+  const done = (async () => {
+    for await (const _m of iter) {
+      // nothing
+    }
+  })().then();
+
+  await deadline(done, 1000);
+
+  await cleanup(ns, nc);
+});

nats-base-client/protocol.ts

@@ -245,6 +245,7 @@ export class SubscriptionImpl extends QueuedIteratorImpl<Msg>
       return Promise.reject(NatsError.errorForCode(ErrorCode.SubClosed));
     }
     if (!this.drained) {
+      this.draining = true;
       this.protocol.unsub(this);
       this.drained = this.protocol.flush(deferred<void>())
         .then(() => {