Rev397-400 on google code; fixes to recycling and minor tweak to heartbeat calling when a disconnection occurs

Lidgren.Network/NetConnection.Handshake.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Collections.Generic;
 using System.Text;
+using System.Threading;
 
 namespace Lidgren.Network
 {
@@ -160,7 +161,7 @@ namespace Lidgren.Network
 			om.Write(m_peerConfiguration.AppIdentifier);
 			om.Write(m_peer.m_uniqueIdentifier);
 			om.Write((float)now);
-
+			Interlocked.Increment(ref om.m_recyclingCount);
 			WriteLocalHail(om);
 
 			if (onLibraryThread)
@@ -184,6 +185,7 @@ namespace Lidgren.Network
 
 			NetOutgoingMessage om = m_peer.CreateMessage(reason);
 			om.m_messageType = NetMessageType.Disconnect;
+			Interlocked.Increment(ref om.m_recyclingCount);
 			if (onLibraryThread)
 				m_peer.SendLibrary(om, m_remoteEndPoint);
 			else

Lidgren.Network/NetConnection.cs

@@ -313,6 +313,8 @@ namespace Lidgren.Network
 				m_sendBufferWritePtr = 0;
 				m_sendBufferNumMessages = 0;
 			}
+
+			Interlocked.Decrement(ref om.m_recyclingCount);
 		}
 
 		/// <summary>
@@ -367,7 +369,6 @@ namespace Lidgren.Network
 				}
 				else
 				{
-
 					switch (method)
 					{
 						case NetDeliveryMethod.Unreliable:

Lidgren.Network/NetNatIntroduction.cs

@@ -1,6 +1,7 @@
 using System;
 using System.Collections.Generic;
 using System.Net;
+using System.Threading;
 
 namespace Lidgren.Network
 {
@@ -23,6 +24,7 @@ namespace Lidgren.Network
 			msg.Write(hostInternal);
 			msg.Write(hostExternal);
 			msg.Write(token);
+			Interlocked.Increment(ref msg.m_recyclingCount);
 			m_unsentUnconnectedMessages.Enqueue(new NetTuple<IPEndPoint, NetOutgoingMessage>(clientExternal, msg));
 
 			// send message to host
@@ -32,6 +34,7 @@ namespace Lidgren.Network
 			msg.Write(clientInternal);
 			msg.Write(clientExternal);
 			msg.Write(token);
+			Interlocked.Increment(ref msg.m_recyclingCount);
 			m_unsentUnconnectedMessages.Enqueue(new NetTuple<IPEndPoint, NetOutgoingMessage>(hostExternal, msg));
 		}
 
@@ -63,6 +66,7 @@ namespace Lidgren.Network
 			punch.m_messageType = NetMessageType.NatPunchMessage;
 			punch.Write(hostByte);
 			punch.Write(token);
+			Interlocked.Increment(ref punch.m_recyclingCount);
 			m_unsentUnconnectedMessages.Enqueue(new NetTuple<IPEndPoint, NetOutgoingMessage>(remoteInternal, punch));
 			LogDebug("NAT punch sent to " + remoteInternal);
 
@@ -71,6 +75,7 @@ namespace Lidgren.Network
 			punch.m_messageType = NetMessageType.NatPunchMessage;
 			punch.Write(hostByte);
 			punch.Write(token);
+			Interlocked.Increment(ref punch.m_recyclingCount);
 			m_unsentUnconnectedMessages.Enqueue(new NetTuple<IPEndPoint, NetOutgoingMessage>(remoteExternal, punch));
 			LogDebug("NAT punch sent to " + remoteExternal);
 
@@ -107,6 +112,7 @@ namespace Lidgren.Network
 			punch.m_messageType = NetMessageType.NatPunchMessage;
 			punch.Write((byte)0);
 			punch.Write(token);
+			Interlocked.Increment(ref punch.m_recyclingCount);
 			m_unsentUnconnectedMessages.Enqueue(new NetTuple<IPEndPoint, NetOutgoingMessage>(senderEndPoint, punch));
 		}
 	}

Lidgren.Network/NetOutgoingMessage.cs

@@ -30,7 +30,7 @@ namespace Lidgren.Network
 	{
 		internal NetMessageType m_messageType;
 		internal bool m_isSent;
-		internal int m_recyclingCount;
+		internal int m_recyclingCount;            // when this reaches zero the message is ready to be recycled
 
 		internal int m_fragmentGroup;             // which group of fragments ths belongs to
 		internal int m_fragmentGroupTotalBits;    // total number of bits in this group
@@ -46,7 +46,7 @@ namespace Lidgren.Network
 			m_messageType = NetMessageType.LibraryError;
 			m_bitLength = 0;
 			m_isSent = false;
-			m_recyclingCount = 0;
+			NetException.Assert(m_recyclingCount == 0);
 			m_fragmentGroup = 0;
 		}
 

Lidgren.Network/NetPeer.Discovery.cs

@@ -1,5 +1,6 @@
 using System;
 using System.Net;
+using System.Threading;
 
 namespace Lidgren.Network
 {
@@ -12,6 +13,7 @@ namespace Lidgren.Network
 		{
 			NetOutgoingMessage om = CreateMessage(0);
 			om.m_messageType = NetMessageType.Discovery;
+			Interlocked.Increment(ref om.m_recyclingCount);
 			m_unsentUnconnectedMessages.Enqueue(new NetTuple<IPEndPoint, NetOutgoingMessage>(new IPEndPoint(IPAddress.Broadcast, serverPort), om));
 		}
 
@@ -34,6 +36,7 @@ namespace Lidgren.Network
 		{
 			NetOutgoingMessage om = CreateMessage(0);
 			om.m_messageType = NetMessageType.Discovery;
+			om.m_recyclingCount = 1;
 			m_unsentUnconnectedMessages.Enqueue(new NetTuple<IPEndPoint, NetOutgoingMessage>(endPoint, om));
 		}
 
@@ -54,6 +57,7 @@ namespace Lidgren.Network
 				throw new NetException("Cannot send discovery message larger than MTU (currently " + m_configuration.MaximumTransmissionUnit + " bytes)");
 
 			msg.m_messageType = NetMessageType.DiscoveryResponse;
+			Interlocked.Increment(ref msg.m_recyclingCount);
 			m_unsentUnconnectedMessages.Enqueue(new NetTuple<IPEndPoint, NetOutgoingMessage>(recipient, msg));
 		}
 	}

Lidgren.Network/NetPeer.Fragmentation.cs

@@ -67,6 +67,8 @@ namespace Lidgren.Network
 				foreach (NetConnection recipient in recipients)
 				{
 					var res = recipient.EnqueueMessage(chunk, method, sequenceChannel);
+					if (res == NetSendResult.Dropped)
+						Interlocked.Decrement(ref chunk.m_recyclingCount);
 					if ((int)res > (int)retval)
 						retval = res; // return "worst" result
 				}

Lidgren.Network/NetPeer.Internal.cs

@@ -340,32 +340,26 @@ namespace Lidgren.Network
 				// do handshake heartbeats
 				if ((m_frameCounter % 3) == 0)
 				{
-					if (m_handshakes.Count > 0)
+					foreach (var kvp in m_handshakes)
 					{
-						lock (m_handshakes)
+						NetConnection conn = kvp.Value as NetConnection;
+#if DEBUG
+						// sanity check
+						if (kvp.Key != kvp.Key)
+							LogWarning("Sanity fail! Connection in handshake list under wrong key!");
+#endif
+						conn.UnconnectedHeartbeat(now);
+						if (conn.m_status == NetConnectionStatus.Connected || conn.m_status == NetConnectionStatus.Disconnected)
 						{
-							foreach (var kvp in m_handshakes)
+#if DEBUG
+							// sanity check
+							if (conn.m_status == NetConnectionStatus.Disconnected && m_handshakes.ContainsKey(conn.RemoteEndPoint))
 							{
-								NetConnection conn = kvp.Value as NetConnection;
+								LogWarning("Sanity fail! Handshakes list contained disconnected connection!");
-#if DEBUG
+								m_handshakes.Remove(conn.RemoteEndPoint);
-								// sanity check
-								if (kvp.Key != kvp.Key)
-									LogWarning("Sanity fail! Connection in handshake list under wrong key!");
-#endif
-								conn.UnconnectedHeartbeat(now);
-								if (conn.m_status == NetConnectionStatus.Connected || conn.m_status == NetConnectionStatus.Disconnected)
-								{
-#if DEBUG
-									// sanity check
-									if (conn.m_status == NetConnectionStatus.Disconnected && m_handshakes.ContainsKey(conn.RemoteEndPoint))
-									{
-										LogWarning("Sanity fail! Handshakes list contained disconnected connection!");
-										m_handshakes.Remove(conn.RemoteEndPoint);
-									}
-#endif
-									break; // collection has been modified
-								}
 							}
+#endif
+							break; // collection has been modified
 						}
 					}
 				}
@@ -384,17 +378,17 @@ namespace Lidgren.Network
 				// do connection heartbeats
 				lock (m_connections)
 				{
-					foreach (NetConnection conn in m_connections)
+					for (int i = m_connections.Count - 1; i >= 0; i--)
 					{
+						var conn = m_connections[i];
 						conn.Heartbeat(now, m_frameCounter);
 						if (conn.m_status == NetConnectionStatus.Disconnected)
 						{
 							//
 							// remove connection
 							//
-							m_connections.Remove(conn);
+							m_connections.RemoveAt(i);
 							m_connectionLookup.Remove(conn.RemoteEndPoint);
-							break; // can't continue iteration here
 						}
 					}
 				}
@@ -406,13 +400,14 @@ namespace Lidgren.Network
 				{
 					NetOutgoingMessage om = unsent.Item2;
 
-					bool connReset;
 					int len = om.Encode(m_sendBuffer, 0, 0);
-					SendPacket(len, unsent.Item1, 1, out connReset);
 
 					Interlocked.Decrement(ref om.m_recyclingCount);
 					if (om.m_recyclingCount <= 0)
 						Recycle(om);
+
+					bool connReset;
+					SendPacket(len, unsent.Item1, 1, out connReset);
 				}
 			}
 
@@ -716,8 +711,7 @@ namespace Lidgren.Network
 					// Ok, start handshake!
 					NetConnection conn = new NetConnection(this, senderEndPoint);
 					conn.m_status = NetConnectionStatus.ReceivedInitiation;
-					lock(m_handshakes)
+					m_handshakes.Add(senderEndPoint, conn);
-						m_handshakes.Add(senderEndPoint, conn);
 					conn.ReceivedHandshake(now, tp, ptr, payloadByteLength);
 					return;
 
@@ -736,11 +730,8 @@ namespace Lidgren.Network
 			// LogDebug("Accepted connection " + conn);
 			conn.InitExpandMTU(NetTime.Now);
 
-			lock (m_handshakes)
+			if (m_handshakes.Remove(conn.m_remoteEndPoint) == false)
-			{
+				LogWarning("AcceptConnection called but m_handshakes did not contain it!");
-				if (m_handshakes.Remove(conn.m_remoteEndPoint) == false)
-					LogWarning("AcceptConnection called but m_handshakes did not contain it!");
-			}
 
 			lock (m_connections)
 			{

Lidgren.Network/NetPeer.MessagePools.cs

@@ -123,6 +123,8 @@ namespace Lidgren.Network
 			if (m_outgoingMessagesPool == null || !m_outgoingMessagesPool.TryDequeue(out retval))
 				retval = new NetOutgoingMessage();
 
+			NetException.Assert(retval.m_recyclingCount == 0, "Wrong recycling count! Should be zero");
+
 			if (initialCapacity > 0)
 				retval.m_data = GetStorage(initialCapacity);
 
@@ -186,6 +188,8 @@ namespace Lidgren.Network
 			if (m_outgoingMessagesPool == null)
 				return;
 
+			NetException.Assert(msg.m_recyclingCount == 0, "Wrong recycling count! Should be zero");
+
 			NetException.Assert(m_outgoingMessagesPool.Contains(msg) == false, "Recyling already recycled message! Thread race?");
 
 			byte[] storage = msg.m_data;

Lidgren.Network/NetPeer.Send.cs

@@ -107,11 +107,10 @@ namespace Lidgren.Network
 				NetException.Assert(sequenceChannel == 0, "Delivery method " + method + " cannot use sequence channels other than 0!");
 			if (msg.m_isSent)
 				throw new NetException("This message has already been sent! Use NetPeer.SendMessage() to send to multiple recipients efficiently");
+			msg.m_isSent = true;
 
 			int mtu = GetMTU(recipients);
 
-			msg.m_isSent = true;
-
 			int len = msg.GetEncodedSize();
 			if (len <= mtu)
 			{
@@ -124,7 +123,7 @@ namespace Lidgren.Network
 						continue;
 					}
 					NetSendResult res = conn.EnqueueMessage(msg, method, sequenceChannel);
-					if (res != NetSendResult.Queued && res != NetSendResult.Sent)
+					if (res == NetSendResult.Dropped)
 						Interlocked.Decrement(ref msg.m_recyclingCount);
 				}
 			}
@@ -151,13 +150,13 @@ namespace Lidgren.Network
 			if (msg.LengthBytes > m_configuration.MaximumTransmissionUnit)
 				throw new NetException("Unconnected messages too long! Must be shorter than NetConfiguration.MaximumTransmissionUnit (currently " + m_configuration.MaximumTransmissionUnit + ")");
 
+			msg.m_isSent = true;
+			msg.m_messageType = NetMessageType.Unconnected;
+
 			IPAddress adr = NetUtility.Resolve(host);
 			if (adr == null)
 				throw new NetException("Failed to resolve " + host);
 
-			msg.m_messageType = NetMessageType.Unconnected;
-			msg.m_isSent = true;
-
 			Interlocked.Increment(ref msg.m_recyclingCount);
 			m_unsentUnconnectedMessages.Enqueue(new NetTuple<IPEndPoint, NetOutgoingMessage>(new IPEndPoint(adr, port), msg));
 		}

Lidgren.Network/NetPeer.cs

@@ -306,8 +306,7 @@ namespace Lidgren.Network
 				conn.m_connectRequested = true;
 				conn.m_connectionInitiator = true;
 
-				lock(m_handshakes)
+				m_handshakes.Add(remoteEndPoint, conn);
-					m_handshakes.Add(remoteEndPoint, conn);
 
 				return conn;
 			}

Lidgren.Network/NetReliableSenderChannel.cs

@@ -139,7 +139,7 @@ namespace Lidgren.Network
 			if (storedMessage != null)
 			{
 #endif
-			if (Interlocked.Decrement(ref storedMessage.m_recyclingCount) <= 0)
+			if (storedMessage.m_recyclingCount <= 0)
 				m_connection.m_peer.Recycle(storedMessage);
 
 #if !DEBUG

Lidgren.Network/NetUnreliableSenderChannel.cs

@@ -48,7 +48,7 @@ namespace Lidgren.Network
 			int left = GetAllowedSends();
 			if (queueLen > left || (message.LengthBytes > m_connection.m_currentMTU && m_connection.m_peerConfiguration.UnreliableSizeBehaviour == NetUnreliableSizeBehaviour.DropAboveMTU))
 			{
-				m_connection.Peer.Recycle(message);
+				// drop message
 				return NetSendResult.Dropped;
 			}
 
@@ -83,7 +83,6 @@ namespace Lidgren.Network
 
 			m_connection.QueueSendMessage(message, seqNr);
 
-			Interlocked.Decrement(ref message.m_recyclingCount);
 			if (message.m_recyclingCount <= 0)
 				m_connection.m_peer.Recycle(message);