Fixed a bug in the channel that was not detecting out of seq messages (meta-pytorch#663)

Summary:

This was because we need to check for out of seq messages
only if it is not the first message of the connection (even if it is
for subsequent sessions).

Differential Revision: D79044526

Author: vidhyav

hyperactor/src/channel/net.rs

@@ -961,13 +961,14 @@ impl<S: AsyncRead + AsyncWrite + Send + 'static + Unpin> ServerConn<S> {
                                 // Ignore retransmits.
                                 Ok(Frame::Message(seq, _)) if seq < next.seq => (),
                                 Ok(Frame::Message(seq, message)) => {
-                                    if seq > next.seq {
+                                    // received seq should be either equal to next seq or
+                                    // it is the first ever payload for this connection (in which case,
+                                    // the other side might be sending a payload from a previous lost connection
+                                    // for the same session). Else error out!
+                                    if seq > next.seq && next.seq != 0 {
                                         tracing::error!("out-of-sequence message from {}", self.source);
-                                        // TODO: T217549393: Enabling
-                                        // this next line causes test
-                                        // `channel::net::tests::test_tcp_reconnect`
-                                        // to unexpectedly fail.
-                                        // break (next, Err(anyhow::anyhow!("out-of-sequence message from {}", self.source)))
+                                        let next_seq = next.seq;
+                                        break (next, Err(anyhow::anyhow!("out-of-sequence message from {}, expected seq {}, got {}", self.source, next_seq, seq)))
                                     }
                                     match tx.send(message).await {
                                         Ok(()) => {
@@ -1173,15 +1174,10 @@ impl SessionManager {
     {
         let session_id = conn.handshake::<M>().await?;
 
-        let session_var = match self.sessions.entry(session_id) {
-            Entry::Occupied(entry) => entry.get().clone(),
-            Entry::Vacant(entry) => {
-                // We haven't seen this session before. We begin with seq=0 and ack=0.
-                let var = MVar::full(Next { seq: 0, ack: 0 });
-                entry.insert(var.clone());
-                var
-            }
-        };
+        let session_var = self
+            .sessions
+            .entry(session_id)
+            .or_insert_with(|| MVar::full(Next { seq: 0, ack: 0 }));
 
         let next = session_var.take().await;
         let (next, res) = conn.process(session_id, tx, cancel_token, next).await;