From 3ba2e19e2b3eb094e162c59f2e7b3428d1e44524 Mon Sep 17 00:00:00 2001
From: Jonas Schnelli <dev@jonasschnelli.ch>
Date: Tue, 29 Mar 2016 14:30:02 +0200
Subject: [PATCH] [ZMQ] append a message sequence number to every ZMQ
 notification

---
 contrib/zmq/zmq_sub.py         | 14 +++++++++-----
 doc/zmq.md                     |  5 +++++
 qa/rpc-tests/zmq_test.py       | 13 +++++++++++++
 src/zmq/zmqpublishnotifier.cpp | 29 +++++++++++++++++++++--------
 src/zmq/zmqpublishnotifier.h   | 12 ++++++++++++
 5 files changed, 60 insertions(+), 13 deletions(-)

diff --git a/contrib/zmq/zmq_sub.py b/contrib/zmq/zmq_sub.py
index decf29d42..6268123dd 100755
--- a/contrib/zmq/zmq_sub.py
+++ b/contrib/zmq/zmq_sub.py
@@ -3,6 +3,7 @@
 import array
 import binascii
 import zmq
+import struct
 
 port = 28332
 
@@ -19,18 +20,21 @@ try:
         msg = zmqSubSocket.recv_multipart()
         topic = str(msg[0])
         body = msg[1]
-
+        sequence = "Unknown";
+        if len(msg[-1]) == 4:
+          msgSequence = struct.unpack('<I', msg[-1])[-1]
+          sequence = str(msgSequence)
         if topic == "hashblock":
-            print "- HASH BLOCK -"
+            print '- HASH BLOCK ('+sequence+') -'
             print binascii.hexlify(body)
         elif topic == "hashtx":
-            print '- HASH TX -'
+            print '- HASH TX  ('+sequence+') -'
             print binascii.hexlify(body)
         elif topic == "rawblock":
-            print "- RAW BLOCK HEADER -"
+            print '- RAW BLOCK HEADER ('+sequence+') -'
             print binascii.hexlify(body[:80])
         elif topic == "rawtx":
-            print '- RAW TX -'
+            print '- RAW TX ('+sequence+') -'
             print binascii.hexlify(body)
 
 except KeyboardInterrupt:
diff --git a/doc/zmq.md b/doc/zmq.md
index 53e557e3d..1e1155349 100644
--- a/doc/zmq.md
+++ b/doc/zmq.md
@@ -99,3 +99,8 @@ using other means such as firewalling.
 Note that when the block chain tip changes, a reorganisation may occur
 and just the tip will be notified. It is up to the subscriber to
 retrieve the chain from the last known block to the new tip.
+
+There are several possibilities that ZMQ notification can get lost
+during transmission depending on the communication type your are
+using. Bitcoind appends an up-counting sequence number to each
+notification which allows listeners to detect lost notifications.
diff --git a/qa/rpc-tests/zmq_test.py b/qa/rpc-tests/zmq_test.py
index 3a8d62ef2..97850bea3 100755
--- a/qa/rpc-tests/zmq_test.py
+++ b/qa/rpc-tests/zmq_test.py
@@ -11,6 +11,7 @@ from test_framework.test_framework import BitcoinTestFramework
 from test_framework.util import *
 import zmq
 import binascii
+import struct
 
 try:
     import http.client as httplib
@@ -47,11 +48,17 @@ class ZMQTest (BitcoinTestFramework):
         print "listen..."
         msg = self.zmqSubSocket.recv_multipart()
         topic = msg[0]
+        assert_equal(topic, b"hashtx")
         body = msg[1]
+        nseq = msg[2]
+        msgSequence = struct.unpack('<I', msg[-1])[-1]
+        assert_equal(msgSequence, 0) #must be sequence 0 on hashtx
 
         msg = self.zmqSubSocket.recv_multipart()
         topic = msg[0]
         body = msg[1]
+        msgSequence = struct.unpack('<I', msg[-1])[-1]
+        assert_equal(msgSequence, 0) #must be sequence 0 on hashblock
         blkhash = bytes_to_hex_str(body)
 
         assert_equal(genhashes[0], blkhash) #blockhash from generate must be equal to the hash received over zmq
@@ -61,12 +68,16 @@ class ZMQTest (BitcoinTestFramework):
         self.sync_all()
 
         zmqHashes = []
+        blockcount = 0
         for x in range(0,n*2):
             msg = self.zmqSubSocket.recv_multipart()
             topic = msg[0]
             body = msg[1]
             if topic == b"hashblock":
                 zmqHashes.append(bytes_to_hex_str(body))
+                msgSequence = struct.unpack('<I', msg[-1])[-1]
+                assert_equal(msgSequence, blockcount+1)
+                blockcount += 1
 
         for x in range(0,n):
             assert_equal(genhashes[x], zmqHashes[x]) #blockhash from generate must be equal to the hash received over zmq
@@ -82,6 +93,8 @@ class ZMQTest (BitcoinTestFramework):
         hashZMQ = ""
         if topic == b"hashtx":
             hashZMQ = bytes_to_hex_str(body)
+            msgSequence = struct.unpack('<I', msg[-1])[-1]
+            assert_equal(msgSequence, blockcount+1)
 
         assert_equal(hashRPC, hashZMQ) #blockhash from generate must be equal to the hash received over zmq
 
diff --git a/src/zmq/zmqpublishnotifier.cpp b/src/zmq/zmqpublishnotifier.cpp
index 707672894..75a2523e7 100644
--- a/src/zmq/zmqpublishnotifier.cpp
+++ b/src/zmq/zmqpublishnotifier.cpp
@@ -122,6 +122,23 @@ void CZMQAbstractPublishNotifier::Shutdown()
     psocket = 0;
 }
 
+bool CZMQAbstractPublishNotifier::SendMessage(const char *command, const void* data, size_t size)
+{
+    assert(psocket);
+
+    /* send three parts, command & data & a LE 4byte sequence number */
+    unsigned char msgseq[sizeof(uint32_t)];
+    WriteLE32(&msgseq[0], nSequence);
+    int rc = zmq_send_multipart(psocket, command, strlen(command), data, size, msgseq, (size_t)sizeof(uint32_t), (void*)0);
+    if (rc == -1)
+        return false;
+
+    /* increment memory only sequence number after sending */
+    nSequence++;
+
+    return true;
+}
+
 bool CZMQPublishHashBlockNotifier::NotifyBlock(const CBlockIndex *pindex)
 {
     uint256 hash = pindex->GetBlockHash();
@@ -129,8 +146,7 @@ bool CZMQPublishHashBlockNotifier::NotifyBlock(const CBlockIndex *pindex)
     char data[32];
     for (unsigned int i = 0; i < 32; i++)
         data[31 - i] = hash.begin()[i];
-    int rc = zmq_send_multipart(psocket, MSG_HASHBLOCK, 9, data, 32, 0);
-    return rc == 0;
+    return SendMessage(MSG_HASHBLOCK, data, 32);
 }
 
 bool CZMQPublishHashTransactionNotifier::NotifyTransaction(const CTransaction &transaction)
@@ -140,8 +156,7 @@ bool CZMQPublishHashTransactionNotifier::NotifyTransaction(const CTransaction &t
     char data[32];
     for (unsigned int i = 0; i < 32; i++)
         data[31 - i] = hash.begin()[i];
-    int rc = zmq_send_multipart(psocket, MSG_HASHTX, 6, data, 32, 0);
-    return rc == 0;
+    return SendMessage(MSG_HASHTX, data, 32);
 }
 
 bool CZMQPublishRawBlockNotifier::NotifyBlock(const CBlockIndex *pindex)
@@ -161,8 +176,7 @@ bool CZMQPublishRawBlockNotifier::NotifyBlock(const CBlockIndex *pindex)
         ss << block;
     }
 
-    int rc = zmq_send_multipart(psocket, MSG_RAWBLOCK, 8, &(*ss.begin()), ss.size(), 0);
-    return rc == 0;
+    return SendMessage(MSG_RAWBLOCK, &(*ss.begin()), ss.size());
 }
 
 bool CZMQPublishRawTransactionNotifier::NotifyTransaction(const CTransaction &transaction)
@@ -171,6 +185,5 @@ bool CZMQPublishRawTransactionNotifier::NotifyTransaction(const CTransaction &tr
     LogPrint("zmq", "zmq: Publish rawtx %s\n", hash.GetHex());
     CDataStream ss(SER_NETWORK, PROTOCOL_VERSION);
     ss << transaction;
-    int rc = zmq_send_multipart(psocket, MSG_RAWTX, 5, &(*ss.begin()), ss.size(), 0);
-    return rc == 0;
+    return SendMessage(MSG_RAWTX, &(*ss.begin()), ss.size());
 }
diff --git a/src/zmq/zmqpublishnotifier.h b/src/zmq/zmqpublishnotifier.h
index 44d5cbea6..22f02a3d0 100644
--- a/src/zmq/zmqpublishnotifier.h
+++ b/src/zmq/zmqpublishnotifier.h
@@ -11,7 +11,19 @@ class CBlockIndex;
 
 class CZMQAbstractPublishNotifier : public CZMQAbstractNotifier
 {
+private:
+    uint32_t nSequence; //! upcounting per message sequence number
+
 public:
+
+    /* send zmq multipart message
+       parts:
+          * command
+          * data
+          * message sequence number
+    */
+    bool SendMessage(const char *command, const void* data, size_t size);
+
     bool Initialize(void *pcontext);
     void Shutdown();
 };