Struggling with server shutdown

[JonasVautherin]

I don't really understand when the tokio::task::spawn_local(rpc_system).await; shuts down. I essentially want to do local_set.await at the end of the program and confirm that all the tasks end.

If I slightly modify the hello_world example to support exactly one connection (i.e. by removing the loop and calling local_set.await), it works:

pub async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = ::std::env::args().collect();
    if args.len() != 3 {
        println!("usage: {} server ADDRESS[:PORT]", args[0]);
        return Ok(());
    }

    let addr = args[2]
        .to_socket_addrs()?
        .next()
        .expect("could not parse address");

    // In my case, I want to share the server between multiple connections so I moved it here
    let hello_world_client: hello_world::Client = capnp_rpc::new_client(HelloWorldImpl);

    // I create a `local_set` object here for later
    let local_set = tokio::task::LocalSet::new();
    local_set.spawn_local(async move {
        let listener = tokio::net::TcpListener::bind(&addr).await?;

        // Note that I removed the loop here
        {
            let (stream, _) = listener.accept().await?;
            stream.set_nodelay(true)?;
            let (reader, writer) =
                tokio_util::compat::TokioAsyncReadCompatExt::compat(stream).split();
            let network = twoparty::VatNetwork::new(
                reader,
                writer,
                rpc_twoparty_capnp::Side::Server,
                Default::default(),
            );

            let rpc_system =
                RpcSystem::new(Box::new(network), Some(hello_world_client.clone().client));

            tokio::task::spawn_local(rpc_system).await;
        }
        Ok::<(), Box<dyn std::error::Error>>(())
    });

    // This awaits until all the tasks are done
    local_set.await;
    Ok(())
}

This does what I expect: the server responds to the first client and shuts down gracefully. Now if I add a client_task in this same process, it just never stops. It's exactly the same as above, but I copied the client.rs part into the same process as a client_task:

pub async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = ::std::env::args().collect();
    if args.len() != 3 {
        println!("usage: {} server ADDRESS[:PORT]", args[0]);
        return Ok(());
    }

    let addr = args[2]
        .to_socket_addrs()?
        .next()
        .expect("could not parse address");

    let hello_world_client: hello_world::Client = capnp_rpc::new_client(HelloWorldImpl);

    let local_set = tokio::task::LocalSet::new();

    // This is the server_task, like above
    local_set.spawn_local(async move {
        let listener = tokio::net::TcpListener::bind(&addr).await?;

        {
            let (stream, _) = listener.accept().await?;
            stream.set_nodelay(true)?;
            let (reader, writer) =
                tokio_util::compat::TokioAsyncReadCompatExt::compat(stream).split();
            let network = twoparty::VatNetwork::new(
                reader,
                writer,
                rpc_twoparty_capnp::Side::Server,
                Default::default(),
            );

            let rpc_system =
                RpcSystem::new(Box::new(network), Some(hello_world_client.clone().client));

            tokio::task::spawn_local(rpc_system).await;
        }
        Ok::<(), Box<dyn std::error::Error>>(())
    });
    
    // This is the client_task, imported from `client.rs`
    local_set.spawn_local(async move {
        let stream = tokio::net::TcpStream::connect(&addr).await?;
        stream.set_nodelay(true)?;
        let (reader, writer) =
            tokio_util::compat::TokioAsyncReadCompatExt::compat(stream).split();
        let rpc_network = Box::new(twoparty::VatNetwork::new(
            reader,
            writer,
            rpc_twoparty_capnp::Side::Client,
            Default::default(),
        ));
        let mut rpc_system = RpcSystem::new(rpc_network, None);
        let hello_world: hello_world::Client =
            rpc_system.bootstrap(rpc_twoparty_capnp::Side::Server);

        tokio::task::spawn_local(rpc_system);

        let mut request = hello_world.say_hello_request();
        request.get().init_request().set_name("blah");

        let reply = request.send().promise.await?;

        println!(
            "received: {}",
            reply.get()?.get_reply()?.get_message()?.to_str()?
        );
        Ok::<(), Box<dyn std::error::Error>>(())
    });

    // This never finishes (?!)
    local_set.await;
    Ok(())
}

This time, local_set.await; never finishes. Though I would have hoped that the client task would finish, free its resources, and then the server task would behave exactly like when client.rs is run in a different process. Obviously that is not the case.

Why is that? What am I missing?

[dwrensha]

When the client and server are in separate processes, the client process terminates and closes the connection, causing the server process to also shut down. When they are in the same process, you need to trigger the disconnect yourself. You can do so using the RpcSystem::get_connector() method:

capnproto-rust/capnp-rpc/src/lib.rs

Lines 295 to 299 in 17a4bec

	/// Returns a `Disconnector` future that can be run to cleanly close the connection to this `RpcSystem`'s network.
	/// You should get the `Disconnector` before you spawn the `RpcSystem`.
	pub fn get_disconnector(&self) -> rpc::Disconnector<VatId> {
	rpc::Disconnector::new(self.connection_state.clone())
	}

This works:

pub async fn main() -> Result<(), Box<dyn std::error::Error>> {
    let args: Vec<String> = ::std::env::args().collect();
    if args.len() != 3 {
        println!("usage: {} server ADDRESS[:PORT]", args[0]);
        return Ok(());
    }

    let addr = args[2]
        .to_socket_addrs()?
        .next()
        .expect("could not parse address");

    let hello_world_client: hello_world::Client = capnp_rpc::new_client(HelloWorldImpl);

    let local_set = tokio::task::LocalSet::new();

    // This is the server_task, like above
    local_set.spawn_local(async move {
        let listener = tokio::net::TcpListener::bind(&addr).await?;

        {
            let (stream, _) = listener.accept().await?;
            stream.set_nodelay(true)?;
            let (reader, writer) =
                tokio_util::compat::TokioAsyncReadCompatExt::compat(stream).split();
            let network = twoparty::VatNetwork::new(
                reader,
                writer,
                rpc_twoparty_capnp::Side::Server,
                Default::default(),
            );

            let rpc_system =
                RpcSystem::new(Box::new(network), Some(hello_world_client.clone().client));

            tokio::task::spawn_local(rpc_system).await;
        }
        Ok::<(), Box<dyn std::error::Error>>(())
    });
    
    // This is the client_task, imported from `client.rs`
    local_set.spawn_local(async move {
        let stream = tokio::net::TcpStream::connect(&addr).await?;
        stream.set_nodelay(true)?;
        let (reader, writer) =
            tokio_util::compat::TokioAsyncReadCompatExt::compat(stream).split();
        let rpc_network = Box::new(twoparty::VatNetwork::new(
            reader,
            writer,
            rpc_twoparty_capnp::Side::Client,
            Default::default(),
        ));
        let mut rpc_system = RpcSystem::new(rpc_network, None);
        let hello_world: hello_world::Client =
            rpc_system.bootstrap(rpc_twoparty_capnp::Side::Server);

        let disconnector = rpc_system.get_disconnector();
        tokio::task::spawn_local(rpc_system);

        let mut request = hello_world.say_hello_request();
        request.get().init_request().set_name("blah");

        let reply = request.send().promise.await?;

        println!(
            "received: {}",
            reply.get()?.get_reply()?.get_message()?.to_str()?
        );
        disconnector.await?;
        Ok::<(), Box<dyn std::error::Error>>(())
    });

    // This finishes.
    local_set.await;
    Ok(())
}

[JonasVautherin]

Ooh! Thanks a lot, that works 🙌!

I actually had tried the disconnector but it was not totally obvious to me that it had to be called from the client (I first tried to implement a stop() function for my server that would call it) and also, I was getting it too early. That is, instead of what you did:

        let mut rpc_system = RpcSystem::new(rpc_network, None);
        let hello_world: hello_world::Client =
            rpc_system.bootstrap(rpc_twoparty_capnp::Side::Server);

        let disconnector = rpc_system.get_disconnector();

I did this:

        let mut rpc_system = RpcSystem::new(rpc_network, None);
        let disconnector = rpc_system.get_disconnector();
        let hello_world: hello_world::Client =
            rpc_system.bootstrap(rpc_twoparty_capnp::Side::Server);

And in that case, disconnector.await has no effect. I see here that the state is set to Disconnected if connection_state is None, and it doesn't get updated later.

At minima, shouldn't it throw instead, letting the user know that they shouldn't get a disconnector before having a connection_state (which comes from the call to bootstrap())?

I do believe that the following patch may fix it, but I am not super familiar with Rust and the Future implementation (but it works for my simple test):

diff --git a/capnp-rpc/src/rpc.rs b/capnp-rpc/src/rpc.rs
index edcb5a60..dd3a3207 100644
--- a/capnp-rpc/src/rpc.rs
+++ b/capnp-rpc/src/rpc.rs
@@ -1642,7 +1642,14 @@ where
                     DisconnectorState::Disconnected
                 }
             }
-            DisconnectorState::Disconnected => DisconnectorState::Disconnected,
+            DisconnectorState::Disconnected => {
+                if self.connection_state.borrow().is_some() {
+                    self.disconnect();
+                    DisconnectorState::Disconnecting
+                } else {
+                    DisconnectorState::Disconnected
+                }
+            },
         };
         match self.state {
             DisconnectorState::Connected => unreachable!(),

My idea is that if the state is "Disconnected" but connection_state is Some, it means that it was initialized like this because it cannot become "Disconnected" if it is Some otherwise.

What do you think? I am happy to open a PR in any case (whether you want to throw or the patch or something else).

[dwrensha]

Good catch! I opened a PR to fix this: #494

[JonasVautherin]

Yes that's even better! Thank you!