[WIP] Add option to merge nats core and jetstream subscriptions

src/js.c

@@ -753,8 +753,94 @@
     natsMsg_Destroy(msg);
 }
 
-static void
-_subComplete(void *closure)
+void
+js_handleSharedReply(natsConnection *nc, natsSubscription *sub, natsMsg *msg, void *closure)
+{
+    const char      *subject    = natsMsg_GetSubject(msg);
+    char            *id         = NULL;
+    jsCtx           *js         = (jsCtx*) closure;
+    natsMsg         *pmsg       = NULL;
+    char            errTxt[256] = {'\0'};
+    jsPubAckErr     pae;
+    jsPubAck        pa;
+    struct jsOptionsPublishAsync *opa = NULL;
+
+    if ((subject == NULL) || (int) strlen(subject) <= js->rpreLen)
+    {
+        return;
+    }
+
+    id = (char*) (subject+js->rpreLen);
+
+    js_lock(js);
+
+    pmsg = natsStrHash_Remove(js->pm, id);
+    if (pmsg == NULL)
+    {
+        js_unlock(js);
+        return;
+    }
+
+    opa = &(js->opts.PublishAsync);
+    if (opa->AckHandler)
+    {
+        jsPubAckErr *ppae   = NULL;
+        jsPubAck    *ppa    = NULL;
+
+        // If _parsePubAck returns an error, we will set the pointer ppae to
+        // our stack variable 'pae', otherwise, set the pointer ppa to the
+        // stack variable 'pa', which is the jsPubAck (positive ack).
+        if (_parsePubAck(msg, &pa, &pae, errTxt, sizeof(errTxt)) != NATS_OK)
+            ppae = &pae;
+        else
+            ppa = &pa;
+
+        // Invoke the handler with pointer to either jsPubAck or jsPubAckErr.
+        js_unlock(js);
+
+        (opa->AckHandler)(js, pmsg, ppa, ppae, opa->AckHandlerClosure);
+
+        js_lock(js);
+
+        _freePubAck(ppa);
+        // Set pmsg to NULL because user was responsible for destroying the message.
+        pmsg = NULL;
+    }
+    else if ((opa->ErrHandler != NULL) && (_parsePubAck(msg, NULL, &pae, errTxt, sizeof(errTxt)) != NATS_OK))
+    {
+        // We will invoke CB only if there is any kind of error.
+        // Associate the message with the pubAckErr object.
+        pae.Msg = pmsg;
+        js_unlock(js);
+
+        (opa->ErrHandler)(js, &pae, opa->ErrHandlerClosure);
+
+        js_lock(js);
+
+        // If the user resent the message, pae->Msg will have been cleared.
+        // In this case, do not destroy the message. Do not blindly destroy
+        // an address that could have been set, so destroy only if pmsg
+        // is same value than pae->Msg.
+        if (pae.Msg != pmsg)
+            pmsg = NULL;
+    }
+
+    // Now that the callback has returned, decrement the number of pending messages.
+    js->pmcount--;
+
+    // If there are callers waiting for async pub completion, or stalled async
+    // publish calls and we are now below max pending, broadcast to unblock them.
+    if (((js->pacw > 0) && (js->pmcount == 0))
+        || ((js->stalled > 0) && (js->pmcount <= opa->MaxPending)))
+    {
+        natsCondition_Broadcast(js->cond);
+    }
+    js_unlock(js);
+    natsMsg_Destroy(pmsg);
+}
+
+void
+js_subComplete(void *closure)
 {
     js_release((jsCtx*) closure);
 }
@@ -801,7 +887,7 @@
             {
                 _retain(js);
                 natsSubscription_SetPendingLimits(js->rsub, -1, -1);
-                natsSubscription_SetOnCompleteCB(js->rsub, _subComplete, (void*) js);
+                natsSubscription_SetOnCompleteCB(js->rsub, js_subComplete, (void*) js);
             }
             NATS_FREE(subj);
         }

src/js.h

@@ -299,3 +299,9 @@ js_maybeFetchMore(natsSubscription *sub, jsFetch *fetch);
 
 void
 js_setOnReleasedCb(jsCtx *js, js_onReleaseCb cb, void *arg);
+
+void
+js_subComplete(void *closure);
+
+void
+js_handleSharedReply(natsConnection *nc, natsSubscription *sub, natsMsg *msg, void *closure);

src/msg.c

@@ -913,6 +913,24 @@
     return NATS_UPDATE_ERR_STACK(s);
 }
 
+natsStatus
+natsMsg_Clone(natsMsg **dest, natsMsg *src)
+{
+    natsMsg    *msg      = NULL;
+
+    if (dest == NULL)
+        return nats_setDefaultError(NATS_INVALID_ARG);
+
+    msg = NATS_MALLOC(sizeof(natsMsg));
+    if (msg == NULL)
+        return nats_setDefaultError(NATS_NO_MEMORY);
+
+    memcpy(msg, src, sizeof(natsMsg));
+    *dest = msg;
+
+    return NATS_OK;
+}
+
 bool
 natsMsg_IsNoResponders(natsMsg *m)
 {

src/msg.h

@@ -129,4 +129,7 @@ natsMsg_freeHeaders(natsMsg *msg);
 void
 natsMsg_free(void *object);
 
+natsStatus
+natsMsg_Clone(natsMsg **dest, natsMsg *src);
+
 #endif /* MSG_H_ */

src/nats.h

@@ -6351,6 +6351,9 @@ natsSubscription_SetOnCompleteCB(natsSubscription *sub, natsOnCompleteCB cb, voi
 NATS_EXTERN void
 natsSubscription_Destroy(natsSubscription *sub);
 
+NATS_EXTERN natsStatus
+natsSubscription_CreateSharedSubscription(natsConnection *nc, jsCtx *js);
+
 /** @} */ // end of subGroup
 
 #if defined(NATS_HAS_STREAMING)

src/natsp.h

@@ -697,6 +697,8 @@ struct __natsSubscription
 
     // For JetStream
     jsSub                       *jsi;
+    // Counts the usage of this subscription when shared between nc and js
+    int                         shareCount;
 };
 
 typedef struct __natsPong

src/sub.c

@@ -1370,10 +1370,158 @@
     if (sub->jsi != NULL)
         sub->jsi->dc = false;
 
+    if (sub->shareCount > 0)
+        sub->shareCount--;
+
     natsSub_Unlock(sub);
 
+    // Don't close or free the subscription if it's still in use
+    if (sub->shareCount != 0)
+        return;
+
     if (doUnsub)
         (void)natsSubscription_Unsubscribe(sub);
 
     natsSub_release(sub);
 }
+
+static void
+_sharedRespHandler(natsConnection *nc, natsSubscription *sub, natsMsg *msg, void *closure)
+{
+    char        *rt   = NULL;
+    const char  *subj = NULL;
+    respInfo    *resp = NULL;
+    bool        dmsg  = true;
+
+    js_handleSharedReply(nc, sub, msg, closure);
+
+    natsConn_Lock(nc);
+    if (natsConn_isClosed(nc))
+    {
+        natsConn_Unlock(nc);
+        natsMsg_Destroy(msg);
+        return;
+    }
+    subj = natsMsg_GetSubject(msg);
+    // We look for the reply token by first checking that the message subject
+    // prefix matches the subscription's subject (without the last '*').
+    // It is possible that it does not due to subject rewrite (JetStream).
+    if (((int) strlen(subj) > nc->reqIdOffset)
+        && (memcmp((const void*) sub->subject, (const void*) subj, strlen(sub->subject) - 1) == 0))
+    {
+        rt = (char*) (natsMsg_GetSubject(msg) + nc->reqIdOffset);
+        resp = (respInfo*) natsStrHash_Remove(nc->respMap, rt);
+    }
+    else if (natsStrHash_Count(nc->respMap) == 1)
+    {
+        // Only if the subject is completely different, we assume that it
+        // could be the server that has rewritten the subject and so if there
+        // is a single entry, use that.
+        void *value = NULL;
+        natsStrHash_RemoveSingle(nc->respMap, NULL, &value);
+        resp = (respInfo*) value;
+    }
+    if (resp != NULL)
+    {
+        natsMutex_Lock(resp->mu);
+        // Check for the race where the requestor has already timed-out.
+        // If so, resp->removed will be true, in which case simply discard
+        // the message.
+        if (!resp->removed)
+        {
+
+            dmsg = false;
+            resp->msg = msg;
+            resp->removed = true;
+            natsCondition_Signal(resp->cond);
+        }
+        natsMutex_Unlock(resp->mu);
+    }
+    natsConn_Unlock(nc);
+
+    if (dmsg)
+        natsMsg_Destroy(msg);
+}
+
+natsStatus
+natsSubscription_CreateSharedSubscription(natsConnection *nc, jsCtx *js)
+{
+    natsSubscription *sub = NULL;
+    natsStatus       s = NATS_OK;
+
+    // If either are already set, we shouldn't create a new one
+    if ((nc->respMux != NULL) || (js->rsub != NULL) || (js->nc != nc)) {
+        return NATS_INVALID_ARG;
+    }
+
+    nc->respPool = NATS_CALLOC(RESP_INFO_POOL_MAX_SIZE, sizeof(respInfo*));
+    if (nc->respPool == NULL)
+        s = nats_setDefaultError(NATS_NO_MEMORY);
+    if (s == NATS_OK)
+        s = natsStrHash_Create(&nc->respMap, 4);
+
+    s = natsCondition_Create(&(js->cond));
+    IFOK(s, natsStrHash_Create(&(js->pm), 64));
+    if (s == NATS_OK)
+    {
+        js->rpre = NATS_MALLOC(js->rpreLen+1);
+        if (js->rpre == NULL)
+            s = nats_setDefaultError(NATS_NO_MEMORY);
+        else
+        {
+            char nuid[NUID_BUFFER_LEN+1];
+
+            s = natsNUID_Next(nuid, sizeof(nuid));
+            if (s == NATS_OK)
+            {
+                memcpy(js->rpre, js->nc->inboxPfx, js->nc->inboxPfxLen);
+                memcpy(js->rpre+js->nc->inboxPfxLen, nuid+((int)strlen(nuid)-8), 8);
+                js->rpre[js->rpreLen-1] = '.';
+                js->rpre[js->rpreLen]   = '\0';
+            }
+        }
+    }
+    if (s == NATS_OK)
+    {
+        char *subj = NULL;
+
+        if (nats_asprintf(&subj, "%s*", js->rpre) < 0)
+            s = nats_setDefaultError(NATS_NO_MEMORY);
+        else
+        {
+            s = natsConn_subscribeNoPool(&sub, js->nc, subj, _sharedRespHandler, (void*) js);
+            if (s == NATS_OK)
+            {
+                nc->respSub = subj;
+                nc->reqIdOffset = js->rpreLen;
+
+                sub->shareCount = 2;
+                nc->respMux = sub;
+                js->rsub = sub;
+            }
+            else {
+                NATS_FREE(subj);
+            }
+        }
+
+        if (s == NATS_OK)
+        {
+            _retain(js);
+            natsSubscription_SetPendingLimits(js->rsub, -1, -1);
+            natsSubscription_SetOnCompleteCB(js->rsub, js_subComplete, (void*) js);
+        }
+    }
+    if (s != NATS_OK)
+    {
+        // Undo the things we created so we retry again next time.
+        // It is either that or we have to always check individual
+        // objects to know if we have to create them.
+        NATS_FREE(js->rpre);
+        js->rpre = NULL;
+        natsStrHash_Destroy(js->pm);
+        js->pm = NULL;
+        natsCondition_Destroy(js->cond);
+        js->cond = NULL;
+    }
+    return NATS_UPDATE_ERR_STACK(s);
+}

test/list_test.txt

@@ -105,6 +105,7 @@ _test(JetStreamPublish)
 _test(JetStreamPublishAckHandler)
 _test(JetStreamPublishAsync)
 _test(JetStreamPublishTTL)
+_test(JetStreamSharedSub)
 _test(JetStreamStreamsSealAndRollup)
 _test(JetStreamSubscribe)
 _test(JetStreamSubscribeConfigCheck)

test/test.c

@@ -42397,3 +42397,57 @@ int main(int argc, char **argv)
     printf("ALL PASSED\n");
     return 0;
 }
+
+void test_JetStreamSharedSub(void)
+{
+    natsStatus          s;
+    natsMsg             *msg      = NULL;
+    natsMsg             *req      = NULL;
+    struct threadArg    arg;
+    natsSubscription    *sub      = NULL;
+
+    JS_SETUP(2, 12, 0);
+
+    s = _createDefaultThreadArgsForCbTests(&arg);
+    if ( s != NATS_OK)
+        FAIL("Unable to setup test!");
+
+    arg.string = "I will help you";
+    arg.status = NATS_OK;
+    arg.control= 4;
+
+    test("Creating shared subscription: ");
+    s = natsSubscription_CreateSharedSubscription(nc, js);
+    testCond(s == NATS_OK);
+
+    test("Subscribe: ");
+    IFOK(s, natsConnection_Subscribe(&sub, nc, "foo", _recvTestString, (void*) &arg));
+    testCond(s == NATS_OK);
+
+    test("Validate Jetstream Publish: ");
+    s = natsMsg_Create(&req, "foo", NULL, "help", 4);
+    IFOK(s, js_PublishMsgAsync(js, &req, NULL));
+    testCond(s == NATS_OK);
+
+    test("Wait for publish ack: ");
+    s = js_PublishAsyncComplete(js, NULL);
+    testCond(s == NATS_OK);
+
+    test("Create req message: ");
+    s = natsMsg_Create(&req, "foo", NULL, "help", 4);
+    testCond(s == NATS_OK);
+
+    test("Validate Core RequestMsg: ");
+    s = natsConnection_RequestMsg(&msg, nc, req, 100);
+    natsMutex_Lock(arg.m);
+    while ((s != NATS_TIMEOUT) && !arg.msgReceived)
+        s = natsCondition_TimedWait(arg.c, arg.m, 2000);
+    testCond(s == NATS_OK);
+    natsMsg_Destroy(req);
+    natsMsg_Destroy(msg);
+
+    natsSubscription_Destroy(sub);
+    _destroyDefaultThreadArgs(&arg);
+
+    JS_TEARDOWN;
+}