Bug 1104156 - TCPSocket doesn't close output stream until all buffered streams in the multiplex stream are sent. r=jdm, r=bkelly

dom/network/TCPSocket.cpp

@@ -776,6 +776,18 @@ TCPSocket::MaybeReportErrorAndCloseIfOpen(nsresult status) {
 
 void
 TCPSocket::Close()
+{
+  CloseHelper(true);
+}
+
+void
+TCPSocket::CloseImmediately()
+{
+  CloseHelper(false);
+}
+
+void
+TCPSocket::CloseHelper(bool waitForUnsentData)
 {
   if (mReadyState == TCPReadyState::Closed || mReadyState == TCPReadyState::Closing) {
     return;
@@ -792,12 +804,13 @@ TCPSocket::Close()
   if (mMultiplexStream) {
     mMultiplexStream->GetCount(&count);
   }
-  if (!count) {
+  if (!count || !waitForUnsentData) {
     if (mSocketOutputStream) {
       mSocketOutputStream->Close();
       mSocketOutputStream = nullptr;
     }
   }
+
   if (mSocketInputStream) {
     mSocketInputStream->Close();
     mSocketInputStream = nullptr;

dom/network/TCPSocket.h

@@ -99,6 +99,7 @@ public:
   void Suspend();
   void Resume(ErrorResult& aRv);
   void Close();
+  void CloseImmediately();
   bool Send(JSContext* aCx, const nsACString& aData, ErrorResult& aRv);
   bool Send(JSContext* aCx,
             const ArrayBuffer& aData,
@@ -186,6 +187,8 @@ private:
   // Helper for FireDataStringEvent/FireDataArrayEvent.
   nsresult FireDataEvent(JSContext* aCx, const nsAString& aType,
                          JS::Handle<JS::Value> aData);
+  // Helper for Close/CloseImmediately
+  void CloseHelper(bool waitForUnsentData);
 
   TCPReadyState mReadyState;
   // Whether to use strings or array buffers for the "data" event.

dom/network/tests/test_tcpsocket_client_and_server_basics.js

@@ -357,6 +357,31 @@ function* test_basics() {
   is((yield clientQueue.waitForEvent()).type, 'close',
      'The close event should fire after the drain event.');
 
+  // -- Re-establish connection (Test for Close Immediately)
+  connectedPromise = waitForConnection(listeningServer);
+  clientSocket = createSocket('127.0.0.1', serverPort,
+                               { binaryType: 'arraybuffer' });
+  clientQueue = listenForEventsOnSocket(clientSocket, 'client');
+  is((yield clientQueue.waitForEvent()).type, 'open', 'got open event');
+
+  connectedResult = yield connectedPromise;
+  // destructuring assignment is not yet ES6 compliant, must manually unpack
+  serverSocket = connectedResult.socket;
+  serverQueue = connectedResult.queue;
+
+  // -- Attempt to send two non-string data.
+  is(clientSocket.send(bigUint8Array.buffer, 0, bigUint8Array.length), false,
+     'Server sending more than 64k should result in the buffer being full.');
+  is(clientSocket.send(bigUint8Array.buffer, 0, bigUint8Array.length), false,
+     'Server sending more than 64k should result in the buffer being full.');
+  clientSocket.closeImmediately();
+
+  serverReceived = yield serverQueue.waitForDataWithAtLeastLength(1);
+
+  is(serverReceived.length < (2 * bigUint8Array.length), true, 'Received array length less than sent array length');
+
+  is((yield serverQueue.waitForEvent()).type, 'close',
+     'The close event is received after calling closeImmediately');
 
   // -- Close the listening server (and try to connect)
   // We want to verify that the server actually closes / stops listening when

dom/webidl/TCPSocket.webidl

@@ -88,6 +88,11 @@ interface TCPSocket : EventTarget {
    */
   void close();
 
+  /**
+   * Close the socket immediately without waiting for unsent data.
+   */
+  [ChromeOnly] void closeImmediately();
+
   /**
    * Write data to the socket.
    *