An easier one shot subscribe API

examples-wasm/examples/subscriptions.rs

@@ -5,7 +5,7 @@
 
 use std::future::IntoFuture;
 
-use graphql_ws_client::{next::ClientBuilder, ws_stream_wasm::Connection};
+use graphql_ws_client::{next::Client, ws_stream_wasm::Connection};
 
 mod schema {
     cynic::use_schema!("../schemas/books.graphql");
@@ -57,7 +57,7 @@ async fn main() {
 
     let connection = Connection::new(ws_conn).await;
 
-    let (mut client, actor) = ClientBuilder::new().build(connection).await.unwrap();
+    let (mut client, actor) = Client::build(connection).await.unwrap();
     wasm_bindgen_futures::spawn_local(actor.into_future());
 
     let mut stream = client.streaming_operation(build_query()).await.unwrap();

src/doc_utils.rs

@@ -0,0 +1,14 @@
+use crate::{next::Message, Error};
+
+pub struct Conn;
+
+#[async_trait::async_trait]
+impl crate::next::Connection for Conn {
+    async fn receive(&mut self) -> Option<Message> {
+        unimplemented!()
+    }
+
+    async fn send(&mut self, _: Message) -> Result<(), Error> {
+        unimplemented!()
+    }
+}

src/lib.rs

@@ -26,6 +26,10 @@ mod client;
 mod logging;
 mod protocol;
 
+#[doc(hidden)]
+#[path = "doc_utils.rs"]
+pub mod __doc_utils;
+
 pub mod graphql;
 pub mod websockets;
 

src/next/builder.rs

@@ -1,29 +1,47 @@
-use std::collections::HashMap;
+use std::{collections::HashMap, future::IntoFuture};
 
-use futures::channel::mpsc;
+use futures::{channel::mpsc, future::BoxFuture, stream::BoxStream, FutureExt, StreamExt};
 use serde::Serialize;
 
-use crate::{logging::trace, protocol::Event, Error};
+use crate::{graphql::GraphqlOperation, logging::trace, protocol::Event, Error};
 
 use super::{
     actor::ConnectionActor,
     connection::{Connection, Message},
-    Client,
+    Client, SubscriptionStream,
 };
 
 /// A websocket client builder
-#[derive(Default)]
 pub struct ClientBuilder {
     payload: Option<serde_json::Value>,
     subscription_buffer_size: Option<usize>,
+    connection: Box<dyn Connection + Send>,
 }
 
-impl ClientBuilder {
-    /// Constructs an AsyncWebsocketClientBuilder
-    pub fn new() -> ClientBuilder {
-        ClientBuilder::default()
+impl super::Client {
+    /// Starts building a new Client.
+    ///
+    /// ```rust
+    ///  use graphql_ws_client::next::{Client};
+    ///  use std::future::IntoFuture;
+    /// # let conn = graphql_ws_client::Conn
+    /// # fn main() -> Result<(), graphql_ws_client_::Error> {
+    /// let (client, actor) = Client::new(conn).await?
+    /// # }
+    /// ```
+    pub fn build<Conn>(connection: Conn) -> ClientBuilder
+    where
+        Conn: Connection + Send + 'static,
+    {
+        ClientBuilder {
+            payload: None,
+            subscription_buffer_size: None,
+            connection: Box::new(connection),
+        }
     }
+}
 
+impl ClientBuilder {
     /// Add payload to `connection_init`
     pub fn payload<NewPayload>(self, payload: NewPayload) -> Result<ClientBuilder, Error>
     where
@@ -44,6 +62,50 @@ impl ClientBuilder {
             ..self
         }
     }
+
+    /// Initialise a Client and use it to run a single streaming operation
+    ///
+    /// ```
+    /// todo!("a doctest")
+    /// ```
+    /// If users want to run mutliple operations on a connection they
+    /// should use the `IntoFuture` impl to construct a `Client`
+    pub async fn streaming_operation<'a, Operation>(
+        self,
+        op: Operation,
+    ) -> Result<SubscriptionStream<Operation>, Error>
+    where
+        Operation: GraphqlOperation + Unpin + Send + 'static,
+    {
+        let (mut client, actor) = self.await?;
+
+        let mut actor_future = actor.into_future().fuse();
+
+        let subscribe_future = client.streaming_operation(op).fuse();
+        futures::pin_mut!(subscribe_future);
+
+        // Temporarily run actor_future while we start the subscription
+        let stream = futures::select! {
+            () = actor_future => {
+                return Err(Error::Unknown("actor ended before subscription started".into()))
+            },
+            result = subscribe_future => {
+                result?
+            }
+        };
+
+        Ok(stream.join(actor_future))
+    }
+}
+
+impl IntoFuture for ClientBuilder {
+    type Output = Result<(Client, ConnectionActor), Error>;
+
+    type IntoFuture = BoxFuture<'static, Self::Output>;
+
+    fn into_future(self) -> Self::IntoFuture {
+        todo!()
+    }
 }
 
 impl ClientBuilder {
@@ -52,18 +114,14 @@ impl ClientBuilder {
     /// Accepts an already built websocket connection, and returns the connection
     /// and a future that must be awaited somewhere - if the future is dropped the
     /// connection will also drop.
-    pub async fn build<Conn>(self, connection: Conn) -> Result<(Client, ConnectionActor), Error>
-    where
-        Conn: Connection + Send + 'static,
-    {
-        self.build_impl(Box::new(connection)).await
-    }
+    pub async fn build(self) -> Result<(Client, ConnectionActor), Error> {
+        let Self {
+            payload,
+            subscription_buffer_size,
+            mut connection,
+        } = self;
 
-    async fn build_impl(
-        self,
-        mut connection: Box<dyn Connection + Send>,
-    ) -> Result<(Client, ConnectionActor), Error> {
-        connection.send(Message::init(self.payload)).await?;
+        connection.send(Message::init(payload)).await?;
 
         // wait for ack before entering receiver loop:
         loop {
@@ -108,7 +166,8 @@ impl ClientBuilder {
 
         let actor = ConnectionActor::new(connection, command_receiver);
 
-        let client = Client::new(command_sender, self.subscription_buffer_size.unwrap_or(5));
+        let client =
+            Client::new_internal(command_sender, self.subscription_buffer_size.unwrap_or(5));
 
         Ok((client, actor))
     }

src/next/mod.rs

@@ -36,7 +36,7 @@ pub struct Client {
 }
 
 impl Client {
-    pub(super) fn new(
+    pub(super) fn new_internal(
         actor: mpsc::Sender<ConnectionCommand>,
         subscription_buffer_size: usize,
     ) -> Self {

src/next/stream.rs

@@ -3,7 +3,12 @@ use std::{
     task::{Context, Poll},
 };
 
-use futures::{channel::mpsc, SinkExt, Stream};
+use futures::{
+    channel::mpsc,
+    future::{self, BoxFuture, Fuse},
+    stream::{self, BoxStream},
+    Future, FutureExt, SinkExt, Stream, StreamExt,
+};
 
 use crate::{graphql::GraphqlOperation, Error};
 
@@ -18,7 +23,7 @@ where
     Operation: GraphqlOperation,
 {
     pub(super) id: usize,
-    pub(super) stream: Pin<Box<dyn Stream<Item = Result<Operation::Response, Error>> + Send>>,
+    pub(super) stream: BoxStream<'static, Result<Operation::Response, Error>>,
     pub(super) actor: mpsc::Sender<ConnectionCommand>,
 }
 
@@ -33,6 +38,16 @@ where
             .await
             .map_err(|error| Error::Send(error.to_string()))
     }
+
+    pub(super) fn join(self, future: Fuse<BoxFuture<'static, ()>>) -> Self
+    where
+        Operation::Response: 'static,
+    {
+        Self {
+            stream: self.stream.join(future).boxed(),
+            ..self
+        }
+    }
 }
 
 impl<Operation> Stream for SubscriptionStream<Operation>
@@ -45,3 +60,59 @@ where
         self.project().stream.as_mut().poll_next(cx)
     }
 }
+
+trait JoinStreamExt<'a> {
+    type Item;
+
+    /// Joins a future onto the execution of a stream returning a stream that also polls
+    /// the given future.
+    ///
+    /// If the future ends the stream will still continue till completion but if the stream
+    /// ends the future will be cancelled.
+    ///
+    /// This can be used when you have the receivng side of a channel and a future that sends
+    /// on that channel - combining the two into a single stream that'll run till the channel
+    /// is exhausted.  If you drop the stream you also cancel the underlying process.
+    fn join(self, future: Fuse<BoxFuture<'static, ()>>) -> impl Stream<Item = Self::Item>;
+}
+
+impl<'a, Item> JoinStreamExt<'a> for BoxStream<'static, Item>
+where
+    Item: 'static,
+{
+    type Item = Item;
+
+    fn join(self, future: Fuse<BoxFuture<'static, ()>>) -> impl Stream<Item = Self::Item> + 'a {
+        futures::stream::unfold(
+            ProducerState::Running(self.fuse(), future),
+            |mut state| async {
+                loop {
+                    match state {
+                        ProducerState::Running(mut stream, mut producer) => {
+                            futures::select! {
+                                output = stream.next() => {
+                                    return Some((output?, ProducerState::Running(stream, producer)));
+                                }
+                                _ = producer => {
+                                    state = ProducerState::Draining(stream);
+                                    continue;
+                                }
+                            }
+                        }
+                        ProducerState::Draining(mut stream) => {
+                            return Some((stream.next().await?, ProducerState::Draining(stream)))
+                        }
+                    }
+                }
+            },
+        )
+    }
+}
+
+enum ProducerState<'a, Item> {
+    Running(
+        stream::Fuse<BoxStream<'a, Item>>,
+        future::Fuse<BoxFuture<'a, ()>>,
+    ),
+    Draining(stream::Fuse<BoxStream<'a, Item>>),
+}

tests/cynic-tests.rs

@@ -75,8 +75,7 @@ async fn main_test() {
 
     println!("Connected");
 
-    let (mut client, actor) = graphql_ws_client::next::ClientBuilder::new()
-        .build(connection)
+    let (mut client, actor) = graphql_ws_client::next::Client::build(connection)
         .await
         .unwrap();