Add more proper cleanup for the server

README.md

@@ -67,9 +67,9 @@ Your handler can now be used by the server. You can easily bind your server to a
 use shared::ExampleServer;
 
 let handler = ExampleHandler { state: 0 };
-let server_task = tokio::spawn(ExampleServer::bind(handler, "127.0.0.1:1234"));
+let server_task = ExampleServer::bind(handler, "127.0.0.1:1234").await;
 // Or, if you're using the 'unix' feature...
-let server_task = tokio::spawn(ExampleServer::bind(handler, "/tmp/sock"));
+let server_task = ExampleServer::bind(handler, "/tmp/sock").await;
 
 ```
 

src/lib.rs

@@ -79,7 +79,7 @@ along with this program.  If not, see <https://www.gnu.org/licenses/>.
 //! To start the server, you simply need to use the generated server struct and
 //! pass it your handler.
 //!
-//! ```no_run
+//! ```rust
 //! # use eagle::Protocol;
 //! # #[derive(Protocol)]
 //! # pub enum Example {
@@ -102,18 +102,17 @@ along with this program.  If not, see <https://www.gnu.org/licenses/>.
 //! # tokio_test::block_on(async {
 //! let handler = Handler;
 //! let address = "127.0.0.1:12345"; // Or, if using the 'unix' feature, "/tmp/eagle.sock"
-//! let server_task = tokio::spawn(ExampleServer::bind(handler, address));
+//! let server = ExampleServer::bind(handler, address).await;
+//! server.close().await;
 //! # });
 //! ```
-//!
-//! Please note the usage of `tokio::spawn`. This is because the `bind` function
-//! will not return until the server is closed. You can use the `abort` method
-//! on the task to close the server.
+//! Once bound, the server will begin listening for incoming connections and
+//! queries. **You must remember to use the `close` method to shut down the server.**
 //!
 //! On the client side, you can simply use the generated client struct to connect
 //! to the server and begin sending queries.
 //!
-//! ```no_run
+//! ```rust
 //! # use eagle::Protocol;
 //! # #[derive(Protocol)]
 //! # pub enum Example {
@@ -136,11 +135,11 @@ along with this program.  If not, see <https://www.gnu.org/licenses/>.
 //! # tokio_test::block_on(async {
 //! # let handler = Handler;
 //! let address = "127.0.0.1:12345"; // Or, if using the 'unix' feature, "/tmp/eagle.sock"
-//! # let server_task = tokio::spawn(ExampleServer::bind(handler, address));
+//! # let server = ExampleServer::bind(handler, address).await;
+//! # tokio::time::sleep(tokio::time::Duration::from_millis(10)).await; // Wait for the server to start
 //! let client = ExampleClient::connect(address).await.unwrap();
-//! # // Wait for the server to start, the developer is responsible for this in production
-//! # tokio::time::sleep(std::time::Duration::from_millis(10)).await;
 //! assert_eq!(client.add(2, 5).await.unwrap(), 7);
+//! # server.close().await;
 //! # });
 //! ```
 //!
@@ -463,6 +462,13 @@ fn derive_protocol(input: proc_macro2::TokenStream) -> proc_macro2::TokenStream
                 sc
             }
 
+            pub async fn close(self) {
+                #info("Closing server");
+                for task in self.tasks.lock().await.drain(..) {
+                    task.abort();
+                }
+            }
+
             pub async fn accept_connections<A: #stream_addr_trait>(
                 &self,
                 addr: A,

tests/full.rs

@@ -73,9 +73,8 @@ async fn e2e() {
     };
     #[cfg(feature = "tcp")]
     let address = format!("127.0.0.1:{}", 10000 + rand::random::<u64>() % 1000);
-    let server_task = tokio::spawn(TestProtocolServer::bind(TrivialServer, address.clone()));
-    // Wait for the server to start, the developer is responsible for this in production
-    tokio::time::sleep(std::time::Duration::from_millis(10)).await;
+    let server = TestProtocolServer::bind(TrivialServer, address.clone()).await;
+    tokio::time::sleep(tokio::time::Duration::from_millis(10)).await; // Wait for the server to start
     let client = TestProtocolClient::connect(address).await.unwrap();
     assert_eq!(client.addition(2, 5).await.unwrap(), 7);
     assert_eq!(
@@ -90,5 +89,5 @@ async fn e2e() {
         "The number is 42"
     );
     client.void().await.unwrap();
-    server_task.abort();
+    server.close().await;
 }