Method handle_master_client

cake-core/src/cake/sharding/worker.rs:299–578

Main loop handling communication with the master.

(
        mut socket: TcpStream,
        client: SocketAddr,
        mut context: WorkerContext<G::Shardable>,
    )

    /// Main loop handling communication with the master.
    async fn handle_master_client(
        mut socket: TcpStream,
        client: SocketAddr,
        mut context: WorkerContext<G::Shardable>,
    ) -> Result<()> {
        // Authenticate if cluster key is set
        if let Some(ref cluster_key) = context.context.args.cluster_key {
            super::auth::authenticate_as_worker(&mut socket, cluster_key)
                .await
                .map_err(|e| anyhow!("[{}] authentication failed: {}", &client, e))?;
            log::debug!("[{}] authenticated", &client);
        }

        // read first message: expect Hello, but handle LayerAssignment for master restarts
        let (latency, _size, first_msg) = Self::read_message_timed(&mut socket).await?;
        match first_msg {
            Message::Hello => { /* normal inference handshake, continue below */ }
            Message::LayerAssignment { ref layers, .. } => {
                // Master restarted and is re-running setup against an already-running worker.
                // Ack the assignment (we already have cached data) and signal ready,
                // then close this connection so the master can reconnect for inference.
                log::info!(
                    "[{}] master re-setup: accepting {} layer assignment(s)",
                    &client,
                    layers.len()
                );
                let ack = Message::LayerAssignmentAck { needs_data: false };
                ack.to_writer(&mut socket).await?;
                Message::WorkerReady.to_writer(&mut socket).await?;
                log::info!("[{}] re-setup complete, closing setup connection", &client);
                return Ok(());
            }
            other => {
                return Err(anyhow!(
                    "[{}] unexpected first message (expected Hello): {:?}",
                    &client,
                    other
                ));
            }
        }

        // send info
        if let Err(e) = Self::write_message_timed(
            &mut socket,
            Message::WorkerInfo(context.to_info(latency.as_millis())),
        )
        .await
        {
            return Err(anyhow!("[{}] could not send worker info: {:?}", &client, e));
        }

        let mut msg_idx = 0;
        let mut avg_ops = 0;
        let mut avg_write = 0;
        let mut avg_read = 0;
        let mut read_buf = Vec::with_capacity(64 * 1024);
        let mut write_buf = Vec::with_capacity(64 * 1024);