Method handle_one

cake-core/tests/protocol.rs:63–112

Handle exactly one client connection, then return.

(self, cluster_key: Option<&str>)

    /// Handle exactly one client connection, then return.
    async fn handle_one(self, cluster_key: Option<&str>) -> anyhow::Result<()> {
        let (mut socket, _client) = self.listener.accept().await?;

        // Auth
        if let Some(key) = cluster_key {
            cake_core::cake::auth::authenticate_as_worker(&mut socket, key).await?;
        }

        // First message
        let (_, first) = Message::from_reader(&mut socket).await?;
        match first {
            Message::Hello => {
                Message::WorkerInfo(mock_worker_info(0))
                    .to_writer(&mut socket)
                    .await?;
            }
            other => return Err(anyhow::anyhow!("unexpected first message: {:?}", other)),
        }

        // Message loop
        loop {
            let result = Message::from_reader(&mut socket).await;
            let (_, msg) = match result {
                Ok(m) => m,
                Err(_) => break, // client disconnected
            };

            match msg {
                Message::SingleOp { x, .. } => {
                    log::debug!("single op");
                    Message::Tensor(x).to_writer(&mut socket).await?;
                }
                Message::Batch { x, .. } => {
                    log::debug!("batch");
                    Message::Tensor(x).to_writer(&mut socket).await?;
                }
                Message::Goodbye => {
                    log::debug!("goodbye");
                    Message::WorkerInfo(mock_worker_info(0))
                        .to_writer(&mut socket)
                        .await?;
                }
                other => {
                    return Err(anyhow::anyhow!("unhandled in mock loop: {:?}", other));
                }
            }
        }

        Ok(())
    }
}

/// Helper: connect to addr, do Hello handshake, return stream + WorkerInfo.