CCBC-1672: handle case when packet gets discarded before IO calls completion handler

When the operation is getting retried, its buffers are getting relocated
to new pipeline, and removed from the old one. But if the write
operation still in the kernel and completion-based IO (e.g. windows
IOCP) is being used, the state of the buffer manager of the old pipeline
have to be adjusted.

Change-Id: I950c519546d7d29a4ce05d02eb57c14e06b2e0e2
Reviewed-on: https://review.couchbase.org/c/libcouchbase/+/233237
Reviewed-by: Jared Casey <jared.casey@couchbase.com>
Tested-by: Build Bot <build@couchbase.com>

Author: avsej

plugins/io/iocp/iocp_iops.c

@@ -370,7 +370,8 @@ static void iops_dtor(lcb_io_opt_t iobase)
     {
         iocp_sockdata_t *sd = LCB_LIST_ITEM(cur, iocp_sockdata_t, list);
 
-        IOCP_LOG(IOCP_WARN, "Leak detected in socket %p (%lu). Refcount=%d", sd, sd->sSocket, sd->refcount);
+        IOCP_LOG(IOCP_WARN, "Leak detected in socket %p (%lu). Refcount=%d", sd, (unsigned long)sd->sSocket,
+                 sd->refcount);
         if (sd->sSocket != INVALID_SOCKET) {
             closesocket(sd->sSocket);
             sd->sSocket = INVALID_SOCKET;

plugins/io/iocp/iocp_timer.c

@@ -67,7 +67,7 @@ iocp_timer_t *iocp_tmq_pop(lcb_list_t *list, lcb_uint64_t now)
 
 void iocp_tmq_add(lcb_list_t *list, iocp_timer_t *timer)
 {
-    IOCP_LOG(IOCP_TRACE, "Adding timer %p with ms %lu", timer, timer->ms);
+    IOCP_LOG(IOCP_TRACE, "Adding timer %p with ms %lu", timer, (unsigned long)timer->ms);
     lcb_list_add_sorted(list, &timer->list, iocp_timer_cmp_asc);
 }
 

src/mc/mcreq-flush-inl.h

@@ -89,11 +89,13 @@ static nb_SIZE mcreq__pktflush_callback(void *p, nb_SIZE hint, void *arg)
  */
 static void mcreq_flush_done_ex(mc_PIPELINE *pl, unsigned nflushed, unsigned expected, lcb_U64 now)
 {
+    unsigned nflushed_extra = 0; // number of bytes flushed by IO, but discarded already
     if (nflushed) {
         mc__FLUSHINFO info = {pl, now};
-        netbuf_end_flush2(&pl->nbmgr, nflushed, mcreq__pktflush_callback, offsetof(mc_PACKET, sl_flushq), &info);
+        nflushed_extra =
+            netbuf_end_flush2(&pl->nbmgr, nflushed, mcreq__pktflush_callback, offsetof(mc_PACKET, sl_flushq), &info);
     }
-    if (nflushed < expected) {
+    if (nflushed < expected || nflushed_extra) {
         netbuf_reset_flush(&pl->nbmgr);
     }
 }

src/mc/mcreq.c

@@ -366,15 +366,16 @@ void mcreq_wipe_packet(mc_PIPELINE *pipeline, mc_PACKET *packet)
 {
     if (!(packet->flags & MCREQ_F_KEY_NOCOPY)) {
         if ((packet->flags & MCREQ_F_DETACHED)) {
-            if (pipeline) {
-                netbuf_cleanup_packet(&pipeline->nbmgr, packet);
-            }
             free(SPAN_BUFFER(&packet->kh_span));
         } else {
             netbuf_mblock_release(&pipeline->nbmgr, &packet->kh_span);
         }
     }
 
+    if (pipeline) {
+        netbuf_cleanup_packet(&pipeline->nbmgr, packet);
+    }
+
     if (!(packet->flags & MCREQ_F_HASVALUE)) {
         return;
     }

src/mcserver/mcserver.cc

@@ -1151,6 +1151,9 @@ void Server::handle_connected(lcbio_SOCKET *sock, lcb_STATUS err, lcbio_OSERR sy
     lcbio_CTXPROCS procs{};
     procs.cb_err = on_error;
     procs.cb_read = on_read;
+
+    netbuf_reset_flush(&nbmgr);
+
     procs.cb_flush_done = on_flush_done;
     procs.cb_flush_ready = on_flush_ready;
     connctx = lcbio_ctx_new(sock, this, &procs, "memcached");

src/netbuf/netbuf.c

@@ -605,7 +605,7 @@ nb_SIZE netbuf_start_flush(nb_MGR *mgr, nb_IOV *iovs, int niov, int *nused)
     return ret;
 }
 
-void netbuf_end_flush(nb_MGR *mgr, unsigned int nflushed)
+unsigned int netbuf_end_flush(nb_MGR *mgr, unsigned int nflushed)
 {
     nb_SENDQ *q = &mgr->sendq;
     sllist_iterator iter;
@@ -635,7 +635,7 @@ void netbuf_end_flush(nb_MGR *mgr, unsigned int nflushed)
             break;
         }
     }
-    lcb_assert(!nflushed);
+    return nflushed;
 }
 
 void netbuf_pdu_enqueue(nb_MGR *mgr, void *pdu, nb_SIZE lloff)
@@ -644,11 +644,17 @@ void netbuf_pdu_enqueue(nb_MGR *mgr, void *pdu, nb_SIZE lloff)
     sllist_append(&q->pdus, (sllist_node *)(void *)((char *)pdu + lloff));
 }
 
-void netbuf_end_flush2(nb_MGR *mgr, unsigned int nflushed, nb_getsize_fn callback, nb_SIZE lloff, void *arg)
+unsigned netbuf_end_flush2(nb_MGR *mgr, unsigned int nflushed, nb_getsize_fn callback, nb_SIZE lloff, void *arg)
 {
     sllist_iterator iter;
     nb_SENDQ *q = &mgr->sendq;
-    netbuf_end_flush(mgr, nflushed);
+
+    unsigned nflushed_extra = netbuf_end_flush(mgr, nflushed);
+
+    /* some operations has been discarded while waiting for IO,
+     * adjusting nflushed to account for the extra */
+    lcb_assert(nflushed_extra <= nflushed);
+    nflushed -= nflushed_extra;
 
     /** Add to the nflushed overflow from last call */
     nflushed += q->pdu_offset;
@@ -672,6 +678,8 @@ void netbuf_end_flush2(nb_MGR *mgr, unsigned int nflushed, nb_getsize_fn callbac
 
     /** Store the remainder of data that wasn't processed for next call */
     q->pdu_offset = nflushed;
+
+    return nflushed_extra;
 }
 
 /******************************************************************************
@@ -741,6 +749,9 @@ void netbuf_cleanup_packet(nb_MGR *mgr, const void *packet)
         if (e->parent == packet) {
             sllist_iter_remove(&mgr->sendq.pending, &iter);
             mblock_release_ptr(&mgr->sendq.elempool, (char *)e, sizeof(*e));
+            if (e == mgr->sendq.last_requested) {
+                netbuf_reset_flush(mgr);
+            }
         }
     }
 }

src/netbuf/netbuf.h

@@ -280,7 +280,7 @@ nb_SIZE netbuf_start_flush(nb_MGR *mgr, nb_IOV *iovs, int niov, int *nused);
  * @param mgr the manager object
  * @param nflushed how much data in bytes was flushed to the network.
  */
-void netbuf_end_flush(nb_MGR *mgr, nb_SIZE nflushed);
+unsigned netbuf_end_flush(nb_MGR *mgr, nb_SIZE nflushed);
 
 /**
  * Reset the flush context for the buffer. This should be called only when the
@@ -405,7 +405,7 @@ void netbuf_pdu_enqueue(nb_MGR *mgr, void *pdu, nb_SIZE lloff);
  */
 typedef nb_SIZE (*nb_getsize_fn)(void *pdu, nb_SIZE remaining, void *arg);
 
-void netbuf_end_flush2(nb_MGR *mgr, unsigned int nflushed, nb_getsize_fn callback, nb_SIZE lloff, void *arg);
+unsigned netbuf_end_flush2(nb_MGR *mgr, unsigned int nflushed, nb_getsize_fn callback, nb_SIZE lloff, void *arg);
 
 /**
  * Ensures that the given internal structures of the manager are not allocated

tests/basic/t_netbuf.cc

@@ -24,6 +24,11 @@
 #endif
 #include "netbuf/netbuf.h"
 
+#include "sllist.h"
+#include "sllist-inl.h"
+
+#include <array>
+
 #define BIG_BUF_SIZE 5000
 #define SMALL_BUF_SIZE 50
 
@@ -456,3 +461,134 @@ TEST_F(NetbufTest, testOutOfOrder)
 
     clean_check(&mgr);
 }
+
+struct my_PACKET {
+    sllist_node slnode{nullptr};
+    nb_SPAN key_{};
+    nb_SPAN value_{};
+    bool is_flushed_{false};
+
+    my_PACKET(nb_MGR *mgr, std::string key, std::string value)
+    {
+        key_.size = key.size();
+        netbuf_mblock_reserve(mgr, &key_);
+
+        value_.size = value.size();
+        netbuf_mblock_reserve(mgr, &value_);
+    }
+
+    void remove_pdu_from(nb_MGR *mgr)
+    {
+        sllist_iterator iter;
+        SLLIST_ITERFOR(&mgr->sendq.pdus, &iter)
+        {
+            my_PACKET *el = SLLIST_ITEM(iter.cur, my_PACKET, slnode);
+            if (el == this) {
+                sllist_iter_remove(&mgr->sendq.pdus, &iter);
+            }
+        }
+    }
+
+    [[nodiscard]] auto key() -> nb_SPAN *
+    {
+        return &key_;
+    }
+
+    [[nodiscard]] auto value() -> nb_SPAN *
+    {
+        return &value_;
+    }
+
+    [[nodiscard]] auto is_flushed() const -> bool
+    {
+        return is_flushed_;
+    }
+
+    void mark_as_flushed()
+    {
+        is_flushed_ = true;
+    }
+
+    [[nodiscard]] auto size() const -> std::size_t
+    {
+        return key_.size + value_.size;
+    }
+};
+
+static nb_SIZE packet_flush_callback(void *pdu, nb_SIZE hint, void * /* arg */)
+{
+    my_PACKET *packet = (my_PACKET *)pdu;
+    if (hint >= packet->size()) {
+        packet->mark_as_flushed();
+    }
+    return packet->size();
+}
+
+TEST_F(NetbufTest, testPacketCleanup)
+{
+    nb_MGR mgr;
+    nb_SETTINGS settings;
+    netbuf_default_settings(&settings);
+    settings.data_basealloc = 1;
+    netbuf_init(&mgr, &settings);
+
+    std::array<my_PACKET, 3> packets{
+        my_PACKET{&mgr, "key_1", "value_1"},
+        my_PACKET{&mgr, "key_2", "value_2"},
+        my_PACKET{&mgr, "key_3", "value_3"},
+    };
+
+    /* enqueue first two packets */
+    for (std::size_t i = 0; i < 2; ++i) {
+        netbuf_enqueue_span(&mgr, packets[i].key(), &packets[i]);
+        netbuf_enqueue_span(&mgr, packets[i].value(), &packets[i]);
+        netbuf_pdu_enqueue(&mgr, &packets[i], offsetof(my_PACKET, slnode));
+    }
+
+    nb_IOV iov[10];
+    nb_SIZE to_flush;
+
+    /* start flushing first two packets */
+    to_flush = netbuf_start_flush(&mgr, iov, 4, NULL);
+    ASSERT_EQ(packets[0].size() + packets[1].size(), to_flush);
+
+    /* discard second packet
+     *
+     * this simulates network failure and relocation of the packet
+     * to some other pipeline while the IO still being processed by
+     * the OS kernel
+     */
+    packets[1].remove_pdu_from(&mgr);
+    netbuf_cleanup_packet(&mgr, &packets[1]);
+
+    /*
+     * OS kernel returned and completed flushing
+     */
+    netbuf_end_flush2(&mgr, to_flush, packet_flush_callback, offsetof(my_PACKET, slnode), NULL);
+    ASSERT_TRUE(packets[0].is_flushed());
+    ASSERT_FALSE(packets[1].is_flushed());
+    ASSERT_FALSE(packets[2].is_flushed());
+
+    /*
+     * enqueue last packet
+     */
+    netbuf_enqueue_span(&mgr, packets[2].key(), &packets[2]);
+    netbuf_enqueue_span(&mgr, packets[2].value(), &packets[2]);
+    netbuf_pdu_enqueue(&mgr, &packets[2], offsetof(my_PACKET, slnode));
+
+    /* start flushing last packet */
+    to_flush = netbuf_start_flush(&mgr, iov, 10, NULL);
+    ASSERT_EQ(packets[2].size(), to_flush);
+
+    netbuf_end_flush2(&mgr, to_flush, packet_flush_callback, 0, NULL);
+    ASSERT_TRUE(packets[0].is_flushed());
+    ASSERT_FALSE(packets[1].is_flushed());
+    ASSERT_TRUE(packets[2].is_flushed());
+
+    for (auto &packet : packets) {
+        netbuf_mblock_release(&mgr, packet.key());
+        netbuf_mblock_release(&mgr, packet.value());
+    }
+
+    clean_check(&mgr);
+}