[Disco] Treat hangup of disco worker process as kShutdown

src/runtime/disco/process_session.cc

@@ -48,11 +48,21 @@ class DiscoPipeMessageQueue : private dmlc::Stream, private DiscoProtocol<DiscoP
   }
 
   TVMArgs Recv() {
-    DequeueNextPacket();
+    bool is_implicit_shutdown = DequeueNextPacket();
     TVMValue* values = nullptr;
     int* type_codes = nullptr;
     int num_args = 0;
-    RPCReference::RecvPackedSeq(&values, &type_codes, &num_args, this);
+
+    if (is_implicit_shutdown) {
+      num_args = 2;
+      values = ArenaAlloc<TVMValue>(num_args);
+      type_codes = ArenaAlloc<int>(num_args);
+      TVMArgsSetter setter(values, type_codes);
+      setter(0, static_cast<int>(DiscoAction::kShutDown));
+      setter(1, 0);
+    } else {
+      RPCReference::RecvPackedSeq(&values, &type_codes, &num_args, this);
+    }
     return TVMArgs(values, type_codes, num_args);
   }
 
@@ -62,18 +72,38 @@ class DiscoPipeMessageQueue : private dmlc::Stream, private DiscoProtocol<DiscoP
     write_buffer_.clear();
   }
 
-  void DequeueNextPacket() {
+  /* \brief Read next packet and reset unpacker
+   *
+   * Read the next packet into `read_buffer_`, releasing all arena
+   * allocations performed by the unpacker and resetting the unpacker
+   * to its initial state.
+   *
+   * \return A boolean value.  If true, this packet should be treated
+   *    equivalently to a `DiscoAction::kShutdown` event.  If false,
+   *    this packet should be unpacked.
+   */
+  bool DequeueNextPacket() {
     uint64_t packet_nbytes = 0;
     int read_size = pipe_.Read(&packet_nbytes, sizeof(packet_nbytes));
+    if (read_size == 0) {
+      // Special case, connection dropped between packets.  Treat as a
+      // request to shutdown.
+      return true;
+    }
+
     ICHECK_EQ(read_size, sizeof(packet_nbytes))
         << "Pipe closed without proper shutdown. Please make sure to explicitly call "
            "`Session::Shutdown`";
     read_buffer_.resize(packet_nbytes);
-    pipe_.Read(read_buffer_.data(), packet_nbytes);
+    read_size = pipe_.Read(read_buffer_.data(), packet_nbytes);
+    ICHECK_EQ(read_size, packet_nbytes)
+        << "Pipe closed without proper shutdown. Please make sure to explicitly call "
+           "`Session::Shutdown`";
     read_offset_ = 0;
     this->RecycleAll();
     RPCCode code = RPCCode::kReturn;
     this->Read(&code);
+    return false;
   }
 
   size_t Read(void* data, size_t size) final {