From 686dd2f56ddce1331e7da8ad8c2f74513cc75867 Mon Sep 17 00:00:00 2001
From: Starbeamrainbowlabs <sbrl@starbeamrainbowlabs.com>
Date: Sun, 19 Feb 2017 16:35:12 +0000
Subject: [PATCH] Put some infrastructure in palce to better handle disconnects
 and shutdowns.

---
 .../Client/Messages/HeartbeatMessage.cs       | 16 ++++++
 .../Client/Messages/IdleDisconnectMessage.cs  | 15 ++++++
 Nibriboard/Client/Messages/Message.cs         |  6 +++
 Nibriboard/Client/Messages/ShutdownMessage.cs | 15 ++++++
 Nibriboard/Client/NibriClient.cs              | 44 ++++++++++++++++-
 Nibriboard/Client/NibriClientManager.cs       | 49 ++++++++++++++++++-
 Nibriboard/ClientFiles/RippleLink.js          | 13 +++++
 Nibriboard/Nibriboard.csproj                  |  3 ++
 Nibriboard/NibriboardServer.cs                |  5 +-
 9 files changed, 163 insertions(+), 3 deletions(-)
 create mode 100644 Nibriboard/Client/Messages/HeartbeatMessage.cs
 create mode 100644 Nibriboard/Client/Messages/IdleDisconnectMessage.cs
 create mode 100644 Nibriboard/Client/Messages/ShutdownMessage.cs

diff --git a/Nibriboard/Client/Messages/HeartbeatMessage.cs b/Nibriboard/Client/Messages/HeartbeatMessage.cs
new file mode 100644
index 0000000..e912f7b
--- /dev/null
+++ b/Nibriboard/Client/Messages/HeartbeatMessage.cs
@@ -0,0 +1,16 @@
+﻿using System;
+
+namespace Nibriboard.Client.Messages
+{
+	/// <summary>
+	/// Represents a heartbeat message. The server sends this periodically to the client,
+	/// to which the client should respond with an identical message.
+	/// </summary>
+	public class HeartbeatMessage : Message
+	{
+		public HeartbeatMessage()
+		{
+		}
+	}
+}
+
diff --git a/Nibriboard/Client/Messages/IdleDisconnectMessage.cs b/Nibriboard/Client/Messages/IdleDisconnectMessage.cs
new file mode 100644
index 0000000..8898a49
--- /dev/null
+++ b/Nibriboard/Client/Messages/IdleDisconnectMessage.cs
@@ -0,0 +1,15 @@
+﻿using System;
+
+namespace Nibriboard.Client.Messages
+{
+	/// <summary>
+	/// Sent by the server to an idle client who's not responding to heartbeats just before the client disconnects them.
+	/// </summary>
+	public class IdleDisconnectMessage : Message
+	{
+		public IdleDisconnectMessage()
+		{
+		}
+	}
+}
+
diff --git a/Nibriboard/Client/Messages/Message.cs b/Nibriboard/Client/Messages/Message.cs
index df7bd30..3df574c 100644
--- a/Nibriboard/Client/Messages/Message.cs
+++ b/Nibriboard/Client/Messages/Message.cs
@@ -4,7 +4,13 @@ namespace Nibriboard.Client.Messages
 {
 	public class Message
 	{
+		/// <summary>
+		/// The name of the event that this message is about.
+		/// </summary>
 		public readonly string Event = string.Empty;
+		/// <summary>
+		/// The date and time that this message was sent.
+		/// </summary>
 		public readonly DateTime TimeSent = DateTime.Now;
 
 		public Message()
diff --git a/Nibriboard/Client/Messages/ShutdownMessage.cs b/Nibriboard/Client/Messages/ShutdownMessage.cs
new file mode 100644
index 0000000..a1da63e
--- /dev/null
+++ b/Nibriboard/Client/Messages/ShutdownMessage.cs
@@ -0,0 +1,15 @@
+﻿using System;
+
+namespace Nibriboard.Client.Messages
+{
+	/// <summary>
+	/// A message that's sent by the server to a client to tell them that the server is shutting down.
+	/// </summary>
+	public class ShutdownMessage : Message
+	{
+		public ShutdownMessage()
+		{
+		}
+	}
+}
+
diff --git a/Nibriboard/Client/NibriClient.cs b/Nibriboard/Client/NibriClient.cs
index 27fb162..a0af2ba 100644
--- a/Nibriboard/Client/NibriClient.cs
+++ b/Nibriboard/Client/NibriClient.cs
@@ -52,10 +52,26 @@ namespace Nibriboard.Client
 		public bool Connected = true;
 		public event NibriDisconnectedEvent Disconnected;
 
+		/// <summary>
+		/// The date and time at which the last message was received from this client.
+		/// </summary>
+		public DateTime LastMessageTime = DateTime.Now;
+		/// <summary>
+		/// The number of milliseconds since we last received a message from this client.
+		/// </summary>
+		/// <value>The milliseconds since last message.</value>
+		public int MillisecondsSinceLastMessage {
+			get {
+				return (int)((DateTime.Now - LastMessageTime).TotalMilliseconds);
+			}
+		}
+
 		/// <summary>
 		/// Whether this client has completed the handshake yet or not.
 		/// </summary>
 		public bool HandshakeCompleted = false;
+
+
 		/// <summary>
 		/// The name this client has assignedto themselves.
 		/// </summary>
@@ -102,18 +118,22 @@ namespace Nibriboard.Client
 			client.ConnectionClosed += (WebSocket socket) => {
 				Connected = false;
 				Disconnected(this);
+				Log.WriteLine("[NibriClient] Client #{0} disconnected.", Id);
 			};
 		}
 
 		private async Task handleMessage(string frame)
 		{
+			// Updatet he last time we received a message from the client
+			LastMessageTime = DateTime.Now;
+
+			// Extract the event name from the message that the client sent.
 			string eventName = JsonUtilities.DeserializeProperty<string>(frame, "Event");
 
 			if(eventName == null) {
 				Log.WriteLine("[NibriClient#{0}] Received message that didn't have an event.", Id);
 				return;
 			}
-
 			if (!messageEventTypes.ContainsKey(eventName)) {
 				Log.WriteLine("[NibriClient#{0}] Received message with invalid event {1}.", Id, eventName);
 				return;
@@ -155,6 +175,28 @@ namespace Nibriboard.Client
 			client.Send(message);
 		}
 
+		/// <summary>
+		/// Sends a heartbeat message to this client.
+		/// </summary>
+		public void SendHeartbeat()
+		{
+			Send(new HeartbeatMessage());
+		}
+
+		/// <summary>
+		/// Closes the connection to the client gracefully.
+		/// </summary>
+		public void CloseConnection(Message lastMessage)
+		{
+			if (!Connected)
+				return;
+			
+			// Tell the client that we're shutting down
+			Send(lastMessage);
+
+			client.Close();
+		}
+
 		/// <summary>
 		/// Generates a new ClientState object representing this client's state at the current time.
 		/// </summary>
diff --git a/Nibriboard/Client/NibriClientManager.cs b/Nibriboard/Client/NibriClientManager.cs
index 0049ee0..c75aaba 100644
--- a/Nibriboard/Client/NibriClientManager.cs
+++ b/Nibriboard/Client/NibriClientManager.cs
@@ -4,6 +4,7 @@ using System.Threading.Tasks;
 using System.Collections.Generic;
 using System.Diagnostics;
 using Nibriboard.Client.Messages;
+using System.Threading;
 
 namespace Nibriboard.Client
 {
@@ -15,6 +16,16 @@ namespace Nibriboard.Client
 		private ClientSettings clientSettings;
 		public List<NibriClient> Clients = new List<NibriClient>();
 
+		/// <summary>
+		/// The cancellation token that's used by the main server to tell us when we should shut down.
+		/// </summary>
+		protected CancellationToken canceller;
+
+		/// <summary>
+		/// The interval at which heatbeats should be sent to the client.
+		/// </summary>
+		public readonly int HeatbeatInterval = 5000;
+
 		/// <summary>
 		/// The number of clients currently connected to this Nibriboard.
 		/// </summary>
@@ -24,9 +35,10 @@ namespace Nibriboard.Client
 			}
 		}
 
-		public NibriClientManager(ClientSettings inClientSettings)
+		public NibriClientManager(ClientSettings inClientSettings, CancellationToken inCancellationToken)
 		{
 			clientSettings = inClientSettings;
+			canceller = inCancellationToken;
 		}
 
 		/// <summary>
@@ -71,6 +83,41 @@ namespace Nibriboard.Client
 			}
 		}
 
+		private async Task ClientMaintenanceMonkey()
+		{
+			while (true) {
+				// Exit if we've been asked to shut down
+				if (canceller.IsCancellationRequested) {
+					close();
+					return;
+				}
+
+				// Disconnect unresponsive clients.
+				foreach (NibriClient client in Clients) {
+					// If we haven't heard from this client in a little while, send a heartbeat message
+					if(client.MillisecondsSinceLastMessage > HeatbeatInterval)
+						client.SendHeartbeat();
+
+					// If the client hasn't sent us a message in a while (even though we sent
+					// them a heartbeat to check on them on the last loop), disconnect them
+					if (client.MillisecondsSinceLastMessage > HeatbeatInterval * 2)
+						client.CloseConnection(new IdleDisconnectMessage());
+				}
+
+				await Task.Delay(HeatbeatInterval);
+			}
+		}
+
+		/// <summary>
+		/// Cleans up this NibriClient manager ready for shutdown.
+		/// </summary>
+		private void close()
+		{
+			// Close the connection to all the remaining nibri clients, telling them that the server is about to shut down
+			foreach (NibriClient client in Clients)
+				client.CloseConnection(new ShutdownMessage());
+		}
+
 		/// <summary>
 		/// Clean up after a client disconnects from the server.
 		/// </summary>
diff --git a/Nibriboard/ClientFiles/RippleLink.js b/Nibriboard/ClientFiles/RippleLink.js
index fd82334..4f95924 100644
--- a/Nibriboard/ClientFiles/RippleLink.js
+++ b/Nibriboard/ClientFiles/RippleLink.js
@@ -20,6 +20,9 @@ class RippleLink extends EventEmitter
 		this.websocket.addEventListener("message", this.handleMessage.bind(this));
 		this.websocket.addEventListener("close", this.handleDisconnection.bind(this));
 		
+		// Respond to heartbeats from the server
+		this.on("Heartbeat", this.handleHeartbeat.bind(this));
+		
 		// Close the socket correctly
 		window.addEventListener("beforeunload", (function(event) {
 			this.websocket.close();
@@ -46,6 +49,16 @@ class RippleLink extends EventEmitter
 		this.emit(message.Event, message);
 	}
 	
+	/**
+	 * Replies to heartbeat messages from the server.
+	 */
+	handleHeartbeat(message) {
+		// Reply with a heartbeat
+		this.send({
+			"Event": "Heartbeat"
+		});
+	}
+	
 	/**
 	 * Sends a message object to the server.
 	 */
diff --git a/Nibriboard/Nibriboard.csproj b/Nibriboard/Nibriboard.csproj
index a04c47f..223b2bc 100644
--- a/Nibriboard/Nibriboard.csproj
+++ b/Nibriboard/Nibriboard.csproj
@@ -82,6 +82,9 @@
     <Compile Include="Utilities\ColourHSL.cs" />
     <Compile Include="Utilities\ToStringJsonConverter.cs" />
     <Compile Include="Client\Messages\HandshakeResponseMessage.cs" />
+    <Compile Include="Client\Messages\ShutdownMessage.cs" />
+    <Compile Include="Client\Messages\IdleDisconnectMessage.cs" />
+    <Compile Include="Client\Messages\HeartbeatMessage.cs" />
   </ItemGroup>
   <ItemGroup>
     <EmbeddedResource Include="ClientFiles\index.html" />
diff --git a/Nibriboard/NibriboardServer.cs b/Nibriboard/NibriboardServer.cs
index 71c01b6..bdb9545 100644
--- a/Nibriboard/NibriboardServer.cs
+++ b/Nibriboard/NibriboardServer.cs
@@ -6,6 +6,7 @@ using IotWeb.Common.Http;
 
 using Nibriboard.RippleSpace;
 using Nibriboard.Client;
+using System.Threading;
 
 namespace Nibriboard
 {
@@ -20,6 +21,8 @@ namespace Nibriboard
 		private ClientSettings clientSettings;
 		private RippleSpaceManager planeManager = new RippleSpaceManager();
 
+		private readonly CancellationTokenSource clientManagerCanceller = new CancellationTokenSource();
+
 		public readonly int Port = 31586;
 
 		public NibriboardServer(int inPort = 31586)
@@ -51,7 +54,7 @@ namespace Nibriboard
 			// Websocket setup
 			httpServer.AddWebSocketRequestHandler(
 				clientSettings.WebsocketPath,
-				new NibriClientManager(clientSettings)
+				new NibriClientManager(clientSettings, clientManagerCanceller.Token)
 			);
 		}
 
