Perform a graceful disconnection if the termination was locally initiated

src/Connection.c

@@ -24,6 +24,7 @@ OPUS_MULTISTREAM_CONFIGURATION HighQualityOpusConfig;
 int OriginalVideoBitrate;
 int AudioPacketDuration;
 bool AudioEncryptionEnabled;
+bool UserRequestedTermination;
 
 // Connection stages
 static const char* stageNames[STAGE_MAX] = {
@@ -55,6 +56,12 @@ void LiInterruptConnection(void) {
 
 // Stop the connection by undoing the step at the current stage and those before it
 void LiStopConnection(void) {
+    // If this was a fully complete connection and we haven't started any termination
+    // logic prior to this point, this termination is user requested.
+    if (stage == STAGE_MAX - 1 && !alreadyTerminated) {
+        UserRequestedTermination = true;
+    }
+
     // Disable termination callbacks now
     alreadyTerminated = true;
 
@@ -191,6 +198,10 @@ int LiStartConnection(PSERVER_INFORMATION serverInfo, PSTREAM_CONFIGURATION stre
     OriginalVideoBitrate = streamConfig->bitrate;
     RemoteAddrString = strdup(serverInfo->address);
 
+    alreadyTerminated = false;
+    ConnectionInterrupted = false;
+    UserRequestedTermination = false;
+    
     // Validate the audio configuration
     if (MAGIC_BYTE_FROM_AUDIO_CONFIG(StreamConfig.audioConfiguration) != 0xCA ||
             CHANNEL_COUNT_FROM_AUDIO_CONFIGURATION(StreamConfig.audioConfiguration) > AUDIO_CONFIGURATION_MAX_CHANNEL_COUNT) {
@@ -245,9 +256,6 @@ int LiStartConnection(PSERVER_INFORMATION serverInfo, PSTREAM_CONFIGURATION stre
         goto Cleanup;
     }
 
-    alreadyTerminated = false;
-    ConnectionInterrupted = false;
-
     Limelog("Initializing platform...");
     ListenerCallbacks.stageStarting(STAGE_PLATFORM_INIT);
     err = initializePlatform();

src/ControlStream.c

@@ -75,6 +75,7 @@ static PPLT_CRYPTO_CONTEXT decryptionCtx;
 #define IDX_TERMINATION 7
 
 #define CONTROL_STREAM_TIMEOUT_SEC 10
+#define CONTROL_STREAM_LINGER_TIMEOUT_SEC 2
 
 static const short packetTypesGen3[] = {
     0x1407, // Request IDR frame
@@ -1053,9 +1054,19 @@ int stopControlStream(void) {
     }
 
     if (peer != NULL) {
-        // We use enet_peer_disconnect_now() so the host knows immediately
-        // of our termination and can cleanup properly for reconnection.
+        if (UserRequestedTermination) {
+            // Gracefully disconnect to ensure the remote host receives all of our final
+            // outbound traffic, including any key up events that might be sent.
+            gracefullyDisconnectEnetPeer(client, peer, CONTROL_STREAM_LINGER_TIMEOUT_SEC * 1000);
+            enet_peer_reset(peer);
+        }
+        else {
+            // This termination was requested by the remote host or caused due to a
+            // connection problem, so just do a quick abortive disconnect. The peer
+            // may be gone by this point.
             enet_peer_disconnect_now(peer, 0);
+        }
+
         peer = NULL;
     }
     if (client != NULL) {

src/Limelight-internal.h

@@ -31,6 +31,7 @@ extern OPUS_MULTISTREAM_CONFIGURATION HighQualityOpusConfig;
 extern int OriginalVideoBitrate;
 extern int AudioPacketDuration;
 extern bool AudioEncryptionEnabled;
+extern bool UserRequestedTermination;
 
 #ifndef UINT24_MAX
 #define UINT24_MAX 0xFFFFFF
@@ -63,6 +64,7 @@ extern bool AudioEncryptionEnabled;
 #define MAGIC_BYTE_FROM_AUDIO_CONFIG(x) ((x) & 0xFF)
 
 int serviceEnetHost(ENetHost* client, ENetEvent* event, enet_uint32 timeoutMs);
+int gracefullyDisconnectEnetPeer(ENetHost* host, ENetPeer* peer, enet_uint32 lingerTimeoutMs);
 int extractVersionQuadFromString(const char* string, int* quad);
 bool isReferenceFrameInvalidationEnabled(void);
 void* extendBuffer(void* ptr, size_t newSize);

src/Misc.c

@@ -6,7 +6,7 @@
 // multiple times for retransmissions to work correctly. It is meant to be a drop-in
 // replacement for enet_host_service(). It also handles cancellation of the connection
 // attempt during the wait.
-int serviceEnetHost(ENetHost* client, ENetEvent* event, enet_uint32 timeoutMs) {
+static int serviceEnetHostInternal(ENetHost* client, ENetEvent* event, enet_uint32 timeoutMs, bool ignoreInterrupts) {
     int ret;
 
     // We need to call enet_host_service() multiple times to make sure retransmissions happen
@@ -14,7 +14,7 @@ int serviceEnetHost(ENetHost* client, ENetEvent* event, enet_uint32 timeoutMs) {
         int selectedTimeout = timeoutMs < ENET_INTERNAL_TIMEOUT_MS ? timeoutMs : ENET_INTERNAL_TIMEOUT_MS;
 
         // We want to report an interrupt event if we are able to read data
-        if (ConnectionInterrupted) {
+        if (!ignoreInterrupts && ConnectionInterrupted) {
             Limelog("ENet wait interrupted\n");
             ret = -1;
             break;
@@ -31,6 +31,46 @@ int serviceEnetHost(ENetHost* client, ENetEvent* event, enet_uint32 timeoutMs) {
     return ret;
 }
 
+int serviceEnetHost(ENetHost* client, ENetEvent* event, enet_uint32 timeoutMs) {
+    return serviceEnetHostInternal(client, event, timeoutMs, false);
+}
+
+// This function performs a graceful disconnect, including lingering until outbound
+// traffic is acked (up until the linger timeout elapses).
+int gracefullyDisconnectEnetPeer(ENetHost* host, ENetPeer* peer, enet_uint32 lingerTimeoutMs) {
+    ENetEvent event;
+    int err;
+
+    // Begin the disconnection process. If this peer is already a zombie, we'll get an
+    // immediate disconnect event that will complete the process. If it's still alive,
+    // we'll get the disconnect event after all outgoing messages have been acked.
+    enet_peer_disconnect_later(peer, 0);
+
+    // We must use the internal function which lets us ignore pending interrupts.
+    while ((err = serviceEnetHostInternal(host, &event, lingerTimeoutMs, true)) > 0) {
+        switch (event.type) {
+        case ENET_EVENT_TYPE_RECEIVE:
+            enet_packet_destroy(event.packet);
+            break;
+        case ENET_EVENT_TYPE_DISCONNECT:
+            Limelog("ENet peer disconnection is complete\n");
+            return 0;
+        default:
+            LC_ASSERT(false);
+            break;
+        }
+    }
+
+    if (err == 0) {
+        Limelog("Timed out waiting for ENet peer to acknowledge disconnection\n");
+    }
+    else {
+        Limelog("Failed to receive ENet peer disconnection acknowledgement\n");
+    }
+
+    return -1;
+}
+
 int extractVersionQuadFromString(const char* string, int* quad) {
     char versionString[128];
     char* nextDot;