SERVER-36885: Make ASIO remember IOCP state when out of resources

3 files changed, 161 insertions, 6 deletions

diff --git a/src/third_party/asio-master/asio/include/asio/detail/impl/win_iocp_io_context.ipp b/src/third_party/asio-master/asio/include/asio/detail/impl/win_iocp_io_context.ipp
index c371b86b250..ebb4b62b179 100644
--- a/src/third_party/asio-master/asio/include/asio/detail/impl/win_iocp_io_context.ipp
+++ b/src/third_party/asio-master/asio/include/asio/detail/impl/win_iocp_io_context.ipp
@@ -252,7 +252,7 @@ void win_iocp_io_context::post_deferred_completion(win_iocp_operation* op)
   op->ready_ = 1;
 
   // Enqueue the operation on the I/O completion port.
-  if (!::PostQueuedCompletionStatus(iocp_.handle, 0, 0, op))
+  if (!::PostQueuedCompletionStatus(iocp_.handle, 0, op->completionKey(), op))
   {
     // Out of resources. Put on completed queue instead.
     mutex::scoped_lock lock(dispatch_mutex_);
@@ -272,7 +272,7 @@ void win_iocp_io_context::post_deferred_completions(
     op->ready_ = 1;
 
     // Enqueue the operation on the I/O completion port.
-    if (!::PostQueuedCompletionStatus(iocp_.handle, 0, 0, op))
+    if (!::PostQueuedCompletionStatus(iocp_.handle, 0, op->completionKey(), op))
     {
       // Out of resources. Put on completed queue instead.
       mutex::scoped_lock lock(dispatch_mutex_);
@@ -298,9 +298,10 @@ void win_iocp_io_context::on_pending(win_iocp_operation* op)
 {
   if (::InterlockedCompareExchange(&op->ready_, 1, 0) == 1)
   {
+    op->completionKey() = overlapped_contains_result;
     // Enqueue the operation on the I/O completion port.
     if (!::PostQueuedCompletionStatus(iocp_.handle,
-          0, overlapped_contains_result, op))
+          0, op->completionKey(), op))
     {
       // Out of resources. Put on completed queue instead.
       mutex::scoped_lock lock(dispatch_mutex_);
@@ -322,9 +323,11 @@ void win_iocp_io_context::on_completion(win_iocp_operation* op,
   op->Offset = last_error;
   op->OffsetHigh = bytes_transferred;
 
+
   // Enqueue the operation on the I/O completion port.
+  op->completionKey() = overlapped_contains_result;
   if (!::PostQueuedCompletionStatus(iocp_.handle,
-        0, overlapped_contains_result, op))
+        0, op->completionKey(), op))
   {
     // Out of resources. Put on completed queue instead.
     mutex::scoped_lock lock(dispatch_mutex_);
@@ -344,9 +347,11 @@ void win_iocp_io_context::on_completion(win_iocp_operation* op,
   op->Offset = ec.value();
   op->OffsetHigh = bytes_transferred;
 
+
   // Enqueue the operation on the I/O completion port.
+  op->completionKey() = overlapped_contains_result;
   if (!::PostQueuedCompletionStatus(iocp_.handle,
-        0, overlapped_contains_result, op))
+        0, op->completionKey(), op))
   {
     // Out of resources. Put on completed queue instead.
     mutex::scoped_lock lock(dispatch_mutex_);
diff --git a/src/third_party/asio-master/asio/include/asio/detail/win_iocp_operation.hpp b/src/third_party/asio-master/asio/include/asio/detail/win_iocp_operation.hpp
index 81d43f07f64..e0cbcc56ced 100644
--- a/src/third_party/asio-master/asio/include/asio/detail/win_iocp_operation.hpp
+++ b/src/third_party/asio-master/asio/include/asio/detail/win_iocp_operation.hpp
@@ -51,6 +51,10 @@ public:
     func_(0, this, asio::error_code(), 0);
   }
 
+  ULONG_PTR& completionKey() {
+      return completionKey_;
+  }
+
 protected:
   typedef void (*func_type)(
       void*, win_iocp_operation*,
@@ -58,7 +62,8 @@ protected:
 
   win_iocp_operation(func_type func)
     : next_(0),
-      func_(func)
+      func_(func),
+      completionKey_(0)
   {
     reset();
   }
@@ -76,6 +81,7 @@ protected:
     OffsetHigh = 0;
     hEvent = 0;
     ready_ = 0;
+    completionKey_ = 0;
   }
 
 private:
@@ -84,6 +90,7 @@ private:
   win_iocp_operation* next_;
   func_type func_;
   long ready_;
+  ULONG_PTR completionKey_;
 };
 
 } // namespace detail
diff --git a/src/third_party/asio-master/patches/0006-MONGO-Fix-IOCP-out-of-resource-handling.patch b/src/third_party/asio-master/patches/0006-MONGO-Fix-IOCP-out-of-resource-handling.patch
new file mode 100644
index 00000000000..af6d78d057e
--- /dev/null
+++ b/src/third_party/asio-master/patches/0006-MONGO-Fix-IOCP-out-of-resource-handling.patch
@@ -0,0 +1,143 @@
+From 28c949880a45147d5cb9d1af16dea7494f68f109 Mon Sep 17 00:00:00 2001
+From: Spencer Jackson <spencer.jackson@mongodb.com>
+Date: Mon, 17 Sep 2018 18:12:33 -0400
+Subject: [PATCH] [PATCH] MONGO: Fix IOCP out of resource handling
+
+---
+ .../asio/detail/impl/win_iocp_io_context.ipp      | 15 ++++++++++-----
+ asio/include/asio/detail/win_iocp_operation.hpp   |  9 ++++++++-
+ 2 files changed, 18 insertions(+), 6 deletions(-)
+
+diff --git a/asio/include/asio/detail/impl/win_iocp_io_context.ipp b/asio/include/asio/detail/impl/win_iocp_io_context.ipp
+index c371b86b..ebb4b62b 100644
+--- a/asio/include/asio/detail/impl/win_iocp_io_context.ipp
++++ b/asio/include/asio/detail/impl/win_iocp_io_context.ipp
+@@ -250,11 +250,11 @@ void win_iocp_io_context::post_deferred_completion(win_iocp_operation* op)
+ {
+   // Flag the operation as ready.
+   op->ready_ = 1;
+ 
+   // Enqueue the operation on the I/O completion port.
+-  if (!::PostQueuedCompletionStatus(iocp_.handle, 0, 0, op))
++  if (!::PostQueuedCompletionStatus(iocp_.handle, 0, op->completionKey(), op))
+   {
+     // Out of resources. Put on completed queue instead.
+     mutex::scoped_lock lock(dispatch_mutex_);
+     completed_ops_.push(op);
+     ::InterlockedExchange(&dispatch_required_, 1);
+@@ -270,11 +270,11 @@ void win_iocp_io_context::post_deferred_completions(
+ 
+     // Flag the operation as ready.
+     op->ready_ = 1;
+ 
+     // Enqueue the operation on the I/O completion port.
+-    if (!::PostQueuedCompletionStatus(iocp_.handle, 0, 0, op))
++    if (!::PostQueuedCompletionStatus(iocp_.handle, 0, op->completionKey(), op))
+     {
+       // Out of resources. Put on completed queue instead.
+       mutex::scoped_lock lock(dispatch_mutex_);
+       completed_ops_.push(op);
+       completed_ops_.push(ops);
+@@ -296,13 +296,14 @@ void win_iocp_io_context::abandon_operations(
+ 
+ void win_iocp_io_context::on_pending(win_iocp_operation* op)
+ {
+   if (::InterlockedCompareExchange(&op->ready_, 1, 0) == 1)
+   {
++    op->completionKey() = overlapped_contains_result;
+     // Enqueue the operation on the I/O completion port.
+     if (!::PostQueuedCompletionStatus(iocp_.handle,
+-          0, overlapped_contains_result, op))
++          0, op->completionKey(), op))
+     {
+       // Out of resources. Put on completed queue instead.
+       mutex::scoped_lock lock(dispatch_mutex_);
+       completed_ops_.push(op);
+       ::InterlockedExchange(&dispatch_required_, 1);
+@@ -320,13 +321,15 @@ void win_iocp_io_context::on_completion(win_iocp_operation* op,
+   op->Internal = reinterpret_cast<ulong_ptr_t>(
+       &asio::error::get_system_category());
+   op->Offset = last_error;
+   op->OffsetHigh = bytes_transferred;
+ 
++
+   // Enqueue the operation on the I/O completion port.
++  op->completionKey() = overlapped_contains_result;
+   if (!::PostQueuedCompletionStatus(iocp_.handle,
+-        0, overlapped_contains_result, op))
++        0, op->completionKey(), op))
+   {
+     // Out of resources. Put on completed queue instead.
+     mutex::scoped_lock lock(dispatch_mutex_);
+     completed_ops_.push(op);
+     ::InterlockedExchange(&dispatch_required_, 1);
+@@ -342,13 +345,15 @@ void win_iocp_io_context::on_completion(win_iocp_operation* op,
+   // Store results in the OVERLAPPED structure.
+   op->Internal = reinterpret_cast<ulong_ptr_t>(&ec.category());
+   op->Offset = ec.value();
+   op->OffsetHigh = bytes_transferred;
+ 
++
+   // Enqueue the operation on the I/O completion port.
++  op->completionKey() = overlapped_contains_result;
+   if (!::PostQueuedCompletionStatus(iocp_.handle,
+-        0, overlapped_contains_result, op))
++        0, op->completionKey(), op))
+   {
+     // Out of resources. Put on completed queue instead.
+     mutex::scoped_lock lock(dispatch_mutex_);
+     completed_ops_.push(op);
+     ::InterlockedExchange(&dispatch_required_, 1);
+diff --git a/asio/include/asio/detail/win_iocp_operation.hpp b/asio/include/asio/detail/win_iocp_operation.hpp
+index 81d43f07..e0cbcc56 100644
+--- a/asio/include/asio/detail/win_iocp_operation.hpp
++++ b/asio/include/asio/detail/win_iocp_operation.hpp
+@@ -49,18 +49,23 @@ public:
+   void destroy()
+   {
+     func_(0, this, asio::error_code(), 0);
+   }
+ 
++  ULONG_PTR& completionKey() {
++      return completionKey_;
++  }
++
+ protected:
+   typedef void (*func_type)(
+       void*, win_iocp_operation*,
+       const asio::error_code&, std::size_t);
+ 
+   win_iocp_operation(func_type func)
+     : next_(0),
+-      func_(func)
++      func_(func),
++      completionKey_(0)
+   {
+     reset();
+   }
+ 
+   // Prevents deletion through this type.
+@@ -74,18 +79,20 @@ protected:
+     InternalHigh = 0;
+     Offset = 0;
+     OffsetHigh = 0;
+     hEvent = 0;
+     ready_ = 0;
++    completionKey_ = 0;
+   }
+ 
+ private:
+   friend class op_queue_access;
+   friend class win_iocp_io_context;
+   win_iocp_operation* next_;
+   func_type func_;
+   long ready_;
++  ULONG_PTR completionKey_;
+ };
+ 
+ } // namespace detail
+ } // namespace asio
+ 
+-- 
+2.19.0
+