2 files changed, 11 insertions, 7 deletions

diff --git a/src/mongo/util/net/message_server.h b/src/mongo/util/net/message_server.h
index 5953ca4cdaa..47c545c4196 100644
--- a/src/mongo/util/net/message_server.h
+++ b/src/mongo/util/net/message_server.h
@@ -43,15 +43,21 @@ public:
     virtual ~MessageHandler() {}
 
     /**
-     * called once when a socket is connected
+     * Called once when a socket is connected.
      */
     virtual void connected(AbstractMessagingPort* p) = 0;
 
     /**
-     * called every time a message comes in
-     * handler is responsible for responding to client
+     * Called every time a message comes in. Handler is responsible for responding to client.
      */
     virtual void process(Message& m, AbstractMessagingPort* p) = 0;
+
+    /**
+     * Called once, either when the client disconnects or when the process is shutting down. After
+     * close() is called, this handler's AbstractMessagingPort pointer (passed in via the
+     * connected() method) is no longer valid.
+     */
+    virtual void close() = 0;
 };
 
 class MessageServer {
diff --git a/src/mongo/util/net/message_server_port.cpp b/src/mongo/util/net/message_server_port.cpp
index 06c1fbc76e5..07a04756151 100644
--- a/src/mongo/util/net/message_server_port.cpp
+++ b/src/mongo/util/net/message_server_port.cpp
@@ -210,6 +210,7 @@ private:
         int64_t counter = 0;
         try {
             handler->connected(portWithHandler.get());
+            ON_BLOCK_EXIT([handler]() { handler->close(); });
 
             while (!inShutdown()) {
                 m.reset();
@@ -222,7 +223,6 @@ private:
                         log() << "end connection " << portWithHandler->psock->remoteString() << " ("
                               << conns << word << " now open)" << endl;
                     }
-                    portWithHandler->shutdown();
                     break;
                 }
 
@@ -238,18 +238,16 @@ private:
         } catch (AssertionException& e) {
             log() << "AssertionException handling request, closing client connection: " << e
                   << endl;
-            portWithHandler->shutdown();
         } catch (SocketException& e) {
             log() << "SocketException handling request, closing client connection: " << e << endl;
-            portWithHandler->shutdown();
         } catch (const DBException&
                      e) {  // must be right above std::exception to avoid catching subclasses
             log() << "DBException handling request, closing client connection: " << e << endl;
-            portWithHandler->shutdown();
         } catch (std::exception& e) {
             error() << "Uncaught std::exception: " << e.what() << ", terminating" << endl;
             dbexit(EXIT_UNCAUGHT);
         }
+        portWithHandler->shutdown();
 
 // Normal disconnect path.
 #ifdef MONGO_CONFIG_SSL