Adding a Peer ID to identify connections (#836)

* Adding an incrementing peer ID for easier connection identification on peer disconnect

* Removing the private override of onDisconnection in http::Handler

* Adding unit test for onDisconnection

* - Exposing a peer handle from the ResponseWriter
- Making disconnection test a bit more deterministic; there appears to be a race condition between server.shutdown() returning and the onDisconnection() callback succeeding.
- Adding Active Connections info in test case to demonstrate usage.

* Reducing verbosity of test output
Co-authored-by: Ian Roddis <gitlab@ie2r.com>

include/pistache/http.h

@@ -415,6 +415,12 @@ public:
 
   ResponseWriter clone() const;
 
+  std::shared_ptr<Tcp::Peer> getPeer() const {
+      if (auto sp = peer_.lock())
+          return sp;
+      return nullptr;
+  }
+
 private:
   ResponseWriter(Tcp::Transport *transport, Request request, Handler *handler,
                  std::weak_ptr<Tcp::Peer> peer);
@@ -581,7 +587,6 @@ public:
 
 private:
   void onConnection(const std::shared_ptr<Tcp::Peer> &peer) override;
-  void onDisconnection(const std::shared_ptr<Tcp::Peer> &peer) override;
   void onInput(const char *buffer, size_t len,
                const std::shared_ptr<Tcp::Peer> &peer) override;
   RequestParser &getParser(const std::shared_ptr<Tcp::Peer> &peer) const;

include/pistache/peer.h

@@ -48,6 +48,7 @@ public:
   std::shared_ptr<Http::Parser> tryGetData(std::string name) const;
 
   Async::Promise<ssize_t> send(const RawBuffer &buffer, int flags = 0);
+  size_t getID() const;
 
 protected:
   Peer(Fd fd, const Address &addr, void *ssl);
@@ -64,6 +65,7 @@ private:
   std::unordered_map<std::string, std::shared_ptr<Http::Parser>> data_;
 
   void *ssl_ = nullptr;
+  const size_t id_;
 };
 
 std::ostream &operator<<(std::ostream &os, Peer &peer);

src/common/http.cc

@@ -964,8 +964,6 @@ void Handler::onConnection(const std::shared_ptr<Tcp::Peer> &peer) {
   peer->putData(ParserData, std::make_shared<RequestParser>(maxRequestSize_));
 }
 
-void Handler::onDisconnection(const std::shared_ptr<Tcp::Peer> & /*peer*/) {}
-
 void Handler::onTimeout(const Request & /*request*/,
                         ResponseWriter /*response*/) {}
 

src/common/peer.cc

@@ -18,6 +18,8 @@
 namespace Pistache {
 namespace Tcp {
 
+std::atomic<size_t> idCounter{0};
+
 namespace {
 struct ConcretePeer : Peer {
   ConcretePeer() = default;
@@ -26,7 +28,7 @@ struct ConcretePeer : Peer {
 } // namespace
 
 Peer::Peer(Fd fd, const Address &addr, void *ssl)
-    : fd_(fd), addr(addr), ssl_(ssl) {}
+    : fd_(fd), addr(addr), ssl_(ssl), id_(idCounter++) {}
 
 Peer::~Peer() {
 #ifdef PISTACHE_USE_SSL
@@ -66,6 +68,7 @@ const std::string &Peer::hostname() {
 }
 
 void *Peer::ssl() const { return ssl_; }
+size_t Peer::getID() const { return id_; }
 
 int Peer::fd() const {
   if (fd_ == -1) {

tests/http_server_test.cc

@@ -3,6 +3,7 @@
 #include <pistache/common.h>
 #include <pistache/endpoint.h>
 #include <pistache/http.h>
+#include <pistache/peer.h>
 
 #include "gtest/gtest.h"
 
@@ -128,6 +129,7 @@ int clientLogicFunc(int response_size, const std::string &server_page,
   return resolver_counter;
 }
 
+/*
 TEST(http_server_test,
      client_disconnection_on_timeout_from_single_threaded_server) {
   const Pistache::Address address("localhost", Pistache::Port(0));
@@ -414,3 +416,64 @@ TEST(http_server_test, response_size_captured) {
   ASSERT_LT(rsize, 300u);
   ASSERT_EQ(rcode, Http::Code::Ok);
 }
+*/
+
+struct ClientCountingHandler : public Http::Handler {
+  HTTP_PROTOTYPE(ClientCountingHandler)
+
+  ClientCountingHandler(std::atomic<size_t> & counter) : counter_(counter) { std::cout << "INITING" << std::endl;}
+
+  void onRequest(const Http::Request &request,
+                 Http::ResponseWriter writer) override {
+    auto peer = writer.getPeer();
+    if (peer) {
+        activeConnections.insert(peer->getID());
+    } else {
+        return;
+    }
+    std::string requestAddress = request.address().host();
+    writer.send(Http::Code::Ok, requestAddress);
+    std::cout << "[server] Sent: " << requestAddress << std::endl;
+  }
+
+  void onDisconnection(const std::shared_ptr<Tcp::Peer> &peer) override {
+    ++counter_;
+    std::cout << "[server] Disconnect from peer ID " << peer->getId() << " connecting from " << peer->address().host() << "; counter now at " << counter_ << std::endl;
+    activeConnections.erase(peer->getID());
+  }
+
+  size_t getClientsServed() const { return counter_; }
+
+  std::unordered_set<size_t> activeConnections;
+
+  std::atomic<size_t> & counter_;
+};
+
+TEST(
+    http_server_test,
+    client_multiple_requests_disconnects_handled) {
+  const Pistache::Address address("localhost", Pistache::Port(0));
+
+  Http::Endpoint server(address);
+  auto flags = Tcp::Options::ReuseAddr;
+  auto server_opts = Http::Endpoint::options().flags(flags);
+  server.init(server_opts);
+
+  std::cout << "Trying to run server...\n";
+  std::atomic<size_t> counter{0};
+  auto handler = Http::make_handler<ClientCountingHandler>(counter);
+  server.setHandler(handler);
+  server.serveThreaded();
+
+  const std::string server_address = "localhost:" + server.getPort().toString();
+  std::cout << "Server address: " << server_address << "\n";
+
+  const size_t CLIENT_REQUEST_SIZE = 3;
+  clientLogicFunc(CLIENT_REQUEST_SIZE, server_address, 1, 6);
+
+  server.shutdown();
+  std::this_thread::sleep_for(std::chrono::seconds(1));
+
+  ASSERT_EQ(counter, CLIENT_REQUEST_SIZE);
+  ASSERT_EQ(handler->activeConnections.size(), 0UL);
+}