Draining Messages Before Disconnect

A feature recently added across the NATS client libraries is the ability to drain connections or subscriptions. Closing a connection, or unsubscribing from a subscription, are generally considered immediate requests. When you close or unsubscribe the library will halt messages in any pending queue or cache for subscribers. When you drain a subscription or connection, it will process any inflight and cached/pending messages before closing.

Drain provides clients that use queue subscriptions with a way to bring down applications without losing any messages. A client can bring up a new queue member, drain and shut down the old queue member, all without losing messages sent to the old client. Without drain, there is the possibility of lost messages due to delivery timing.

The libraries can provide drain on a connection or on a subscriber, or both.

For a connection the process is essentially:

  1. Drain all subscriptions
  2. Stop new messages from being published
  3. Flush any remaining published messages
  4. Close

The API for drain can generally be used instead of close:

As an example of draining a connection:

Go

  1. wg := sync.WaitGroup{}
  2. wg.Add(1)
  4. errCh := make(chan error, 1)
  6. // To simulate a timeout, you would set the DrainTimeout()
  7. // to a value less than the time spent in the message callback,
  8. // so say: nats.DrainTimeout(10*time.Millisecond).
  10. nc, err := nats.Connect("demo.nats.io",
  11.     nats.DrainTimeout(10*time.Second),
  12.     nats.ErrorHandler(func(_ *nats.Conn, _ *nats.Subscription, err error) {
  13.         errCh <- err
  14.     }),
  15.     nats.ClosedHandler(func(_ *nats.Conn) {
  16.         wg.Done()
  17.     }))
  18. if err != nil {
  19.     log.Fatal(err)
  20. }
  22. // Just to not collide using the demo server with other users.
  23. subject := nats.NewInbox()
  25. // Subscribe, but add some delay while processing.
  26. if _, err := nc.Subscribe(subject, func(_ *nats.Msg) {
  27.     time.Sleep(200 * time.Millisecond)
  28. }); err != nil {
  29.     log.Fatal(err)
  30. }
  32. // Publish a message
  33. if err := nc.Publish(subject, []byte("hello")); err != nil {
  34.     log.Fatal(err)
  35. }
  37. // Drain the connection, which will close it when done.
  38. if err := nc.Drain(); err != nil {
  39.     log.Fatal(err)
  40. }
  42. // Wait for the connection to be closed.
  43. wg.Wait()
  45. // Check if there was an error
  46. select {
  47. case e := <-errCh:
  48.     log.Fatal(e)
  49. default:
  50. }

Java

  1. Connection nc = Nats.connect("nats://demo.nats.io:4222");
  3. // Use a latch to wait for a message to arrive
  4. CountDownLatch latch = new CountDownLatch(1);
  6. // Create a dispatcher and inline message handler
  7. Dispatcher d = nc.createDispatcher((msg) -> {
  8.     String str = new String(msg.getData(), StandardCharsets.UTF_8);
  9.     System.out.println(str);
  10.     latch.countDown();
  11. });
  13. // Subscribe
  14. d.subscribe("updates");
  16. // Wait for a message to come in
  17. latch.await();
  19. // Drain the connection, which will close it
  20. CompletableFuture<Boolean> drained = nc.drain(Duration.ofSeconds(10));
  22. // Wait for the drain to complete
  23. drained.get();

JavaScript

  1. let nc = NATS.connect({url: "nats://demo.nats.io:4222"});
  2. let inbox = createInbox();
  3. let counter = 0;
  4. nc.subscribe(inbox, () => {
  5.     counter++;
  6. });
  8. nc.publish(inbox);
  9. nc.drain((err)=> {
  10.     if(err) {
  11.         t.log(err);
  12.     }
  13.     t.log('connection is closed:', nc.closed);
  14.     t.log('processed', counter, 'messages');
  15.     t.pass();
  16.     // the snippet is running as a promise in a test
  17.     // and calls resolve to pass the test
  18.     resolve();
  19. });

Python

  1. import asyncio
  2. from nats.aio.client import Client as NATS
  4. async def example(loop):
  5.     nc = NATS()
  7.     await nc.connect("nats://127.0.0.1:4222", loop=loop)
  9.     async def handler(msg):
  10.         print("[Received] ", msg)
  11.         await nc.publish(msg.reply, b'I can help')
  13.         # Can check whether client is in draining state
  14.         if nc.is_draining:
  15.             print("Connection is draining")
  17.     await nc.subscribe("help", "workers", cb=handler)
  18.     await nc.flush()
  20.     requests = []
  21.     for i in range(0, 10):
  22.         request = nc.request("help", b'help!', timeout=1)
  23.         requests.append(request)
  25.     # Wait for all the responses
  26.     responses = []
  27.     responses = await asyncio.gather(*requests)
  29.     # Gracefully close the connection.
  30.     await nc.drain()
  32.     print("Received {} responses".format(len(responses)))

Ruby

  1. NATS.start(drain_timeout: 1) do |nc|
  2.   NATS.subscribe('foo', queue: "workers") do |msg, reply, sub|
  3.     nc.publish(reply, "ACK:#{msg}")
  4.   end
  6.   NATS.subscribe('bar', queue: "workers") do |msg, reply, sub|
  7.     nc.publish(reply, "ACK:#{msg}")
  8.   end
  10.   NATS.subscribe('quux', queue: "workers") do |msg, reply, sub|
  11.     nc.publish(reply, "ACK:#{msg}")
  12.   end
  14.   EM.add_timer(2) do
  15.     next if NATS.draining?
  17.     # Drain gracefully closes the connection.
  18.     NATS.drain do
  19.       puts "Done draining. Connection is closed."
  20.     end
  21.   end
  22. end

TypeScript

  1. let sub = await nc.subscribe('updates', (err, msg) => {
  2.     t.log('worker got message', msg.data);
  3. }, {queue: "workers"});
  4. // [end drain_sub]
  5. nc.flush();
  7. await nc.drain();
  8. // client must close when the connection drain resolves
  9. nc.close();

C

  1. static void
  2. onMsg(natsConnection *conn, natsSubscription *sub, natsMsg *msg, void *closure)
  3. {
  4.     printf("Received msg: %s - %.*s\n",
  5.            natsMsg_GetSubject(msg),
  6.            natsMsg_GetDataLength(msg),
  7.            natsMsg_GetData(msg));
  9.     // Add some delay while processing
  10.     nats_Sleep(200);
  12.     // Need to destroy the message!
  13.     natsMsg_Destroy(msg);
  14. }
  16. static void
  17. closeHandler(natsConnection *conn, void *closure)
  18. {
  19.     cond_variable cv = (cond_variable) closure;
  21.     notify_cond_variable(cv);
  22. }
  24. (...)
  27. natsConnection      *conn      = NULL;
  28. natsOptions         *opts      = NULL;
  29. natsSubscription    *sub       = NULL;
  30. natsStatus          s          = NATS_OK;
  31. cond_variable       cv         = new_cond_variable(); // some fictuous way to notify between threads.
  33. s = natsOptions_Create(&opts);
  34. if (s == NATS_OK)
  35.     // Setup a close handler and pass a reference to our condition variable.
  36.     s = natsOptions_SetClosedCB(opts, closeHandler, (void*) cv);
  37. if (s == NATS_OK)
  38.     s = natsConnection_Connect(&conn, opts);
  40. // Subscribe
  41. if (s == NATS_OK)
  42.     s = natsConnection_Subscribe(&sub, conn, "foo", onMsg, NULL);
  44. // Publish a message
  45. if (s == NATS_OK)
  46.     s = natsConnection_PublishString(conn, "foo", "hello");
  48. // Drain the connection, which will close it when done.
  49. if (s == NATS_OK)
  50.     s = natsConnection_Drain(conn);
  52. // Wait for the connection to be closed
  53. if (s == NATS_OK)
  54.     cond_variable_wait(cv);
  56. (...)
  58. // Destroy objects that were created
  59. natsSubscription_Destroy(sub);
  60. natsConnection_Destroy(conn);
  61. natsOptions_Destroy(opts);

The mechanics of drain for a subscription are simpler:

  1. Unsubscribe
  2. Process all cached or inflight messages
  3. Clean up

The API for drain can generally be used instead of unsubscribe:

Go

  1.     nc, err := nats.Connect("demo.nats.io")
  2.     if err != nil {
  3.         log.Fatal(err)
  4.     }
  5.     defer nc.Close()
  7.     done := sync.WaitGroup{}
  8.     done.Add(1)
  10.     count := 0
  11.     errCh := make(chan error, 1)
  13.     msgAfterDrain := "not this one"
  15.     // Just to not collide using the demo server with other users.
  16.     subject := nats.NewInbox()
  18.     // This callback will process each message slowly
  19.     sub, err := nc.Subscribe(subject, func(m *nats.Msg) {
  20.         if string(m.Data) == msgAfterDrain {
  21.             errCh <- fmt.Errorf("Should not have received this message")
  22.             return
  23.         }
  24.         time.Sleep(100 * time.Millisecond)
  25.         count++
  26.         if count == 2 {
  27.             done.Done()
  28.         }
  29.     })
  31.     // Send 2 messages
  32.     for i := 0; i < 2; i++ {
  33.         nc.Publish(subject, []byte("hello"))
  34.     }
  36.     // Call Drain on the subscription. It unsubscribes but
  37.     // wait for all pending messages to be processed.
  38.     if err := sub.Drain(); err != nil {
  39.         log.Fatal(err)
  40.     }
  42.     // Send one more message, this message should not be received
  43.     nc.Publish(subject, []byte(msgAfterDrain))
  45.     // Wait for the subscription to have processed the 2 messages.
  46.     done.Wait()
  48.     // Now check that the 3rd message was not received
  49.     select {
  50.     case e := <-errCh:
  51.         log.Fatal(e)
  52.     case <-time.After(200 * time.Millisecond):
  53.         // OK!
  54.     }

Java

  1. Connection nc = Nats.connect("nats://demo.nats.io:4222");
  3. // Use a latch to wait for a message to arrive
  4. CountDownLatch latch = new CountDownLatch(1);
  6. // Create a dispatcher and inline message handler
  7. Dispatcher d = nc.createDispatcher((msg) -> {
  8.     String str = new String(msg.getData(), StandardCharsets.UTF_8);
  9.     System.out.println(str);
  10.     latch.countDown();
  11. });
  13. // Subscribe
  14. d.subscribe("updates");
  16. // Wait for a message to come in
  17. latch.await();
  19. // Messages that have arrived will be processed
  20. CompletableFuture<Boolean> drained = d.drain(Duration.ofSeconds(10));
  22. // Wait for the drain to complete
  23. drained.get();
  25. // Close the connection
  26. nc.close();

JavaScript

  1. let nc = NATS.connect({url: "nats://demo.nats.io:4222"});
  2. let inbox = createInbox();
  3. let counter = 0;
  4. let sid = nc.subscribe(inbox, () => {
  5.     counter++;
  6. });
  8. nc.publish(inbox);
  9. nc.drainSubscription(sid, (err)=> {
  10.     if(err) {
  11.         t.log(err);
  12.     }
  13.     t.log('processed', counter, 'messages');
  14. });
  15. nc.flush(() => {
  16.     nc.close();
  17.     t.pass();
  18.     resolve();
  19. });

Python

  1. import asyncio
  2. from nats.aio.client import Client as NATS
  4. async def example(loop):
  5.     nc = NATS()
  7.     await nc.connect("nats://127.0.0.1:4222", loop=loop)
  9.     async def handler(msg):
  10.         print("[Received] ", msg)
  11.         await nc.publish(msg.reply, b'I can help')
  13.         # Can check whether client is in draining state
  14.         if nc.is_draining:
  15.             print("Connection is draining")
  17.     sid = await nc.subscribe("help", "workers", cb=handler)
  18.     await nc.flush()
  20.     # Gracefully unsubscribe the subscription
  21.     await nc.drain(sid)

Ruby

  1. # There is currently no API to drain a single subscription, the whole connection can be drained though via NATS.drain

TypeScript

  1. let sub = await nc.subscribe('updates', (err, msg) => {
  2.     t.log('worker got message', msg.data);
  3. }, {queue: "workers"});

C

  1. natsConnection      *conn      = NULL;
  2. natsSubscription    *sub       = NULL;
  3. natsStatus          s          = NATS_OK;
  5. s = natsConnection_ConnectTo(&conn, NATS_DEFAULT_URL);
  7. // Subscribe
  8. if (s == NATS_OK)
  9.     s = natsConnection_Subscribe(&sub, conn, "foo", onMsg, NULL);
  11. // Publish 2 messages
  12. if (s == NATS_OK)
  13. {
  14.     int i;
  15.     for (i=0; (s == NATS_OK) && (i<2); i++)
  16.     {
  17.         s = natsConnection_PublishString(conn, "foo", "hello");
  18.     }
  19. }
  21. // Call Drain on the subscription. It unsubscribes but
  22. // wait for all pending messages to be processed.
  23. if (s == NATS_OK)
  24.     s = natsSubscription_Drain(sub);
  26. (...)
  28. // Destroy objects that were created
  29. natsSubscription_Destroy(sub);
  30. natsConnection_Destroy(conn);

Because draining can involve messages flowing to the server, for a flush and asynchronous message processing, the timeout for drain should generally be higher than the timeout for a simple message request/reply or similar.