tenant.openagents/omega

1use super::{
2    Connection,
3    message_stream::{Message, MessageStream},
4    proto::{
5        self, AnyTypedEnvelope, EnvelopedMessage, PeerId, Receipt, RequestMessage, TypedEnvelope,
6    },
7};
8use anyhow::{Context as _, Result, anyhow};
9use collections::HashMap;
10use futures::{
11    FutureExt, SinkExt, StreamExt, TryFutureExt,
12    channel::{mpsc, oneshot},
13    stream::BoxStream,
14};
15use parking_lot::{Mutex, RwLock};
16use proto::{ErrorCode, ErrorCodeExt, ErrorExt, RpcError};
17use serde::{Serialize, ser::SerializeStruct};
18use std::{
19    fmt, future,
20    future::Future,
21    sync::atomic::Ordering::SeqCst,
22    sync::{
23        Arc,
24        atomic::{self, AtomicU32},
25    },
26    time::Duration,
27    time::Instant,
28};
29
30#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, Debug, Serialize)]
31pub struct ConnectionId {
32    pub owner_id: u32,
33    pub id: u32,
34}
35
36impl From<ConnectionId> for PeerId {
37    fn from(id: ConnectionId) -> Self {
38        PeerId {
39            owner_id: id.owner_id,
40            id: id.id,
41        }
42    }
43}
44
45impl From<PeerId> for ConnectionId {
46    fn from(peer_id: PeerId) -> Self {
47        Self {
48            owner_id: peer_id.owner_id,
49            id: peer_id.id,
50        }
51    }
52}
53
54impl fmt::Display for ConnectionId {
55    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
56        write!(f, "{}/{}", self.owner_id, self.id)
57    }
58}
59
60pub struct Peer {
61    epoch: AtomicU32,
62    pub connections: RwLock<HashMap<ConnectionId, ConnectionState>>,
63    next_connection_id: AtomicU32,
64}
65
66#[derive(Clone, Serialize)]
67pub struct ConnectionState {
68    #[serde(skip)]
69    outgoing_tx: mpsc::UnboundedSender<Message>,
70    next_message_id: Arc<AtomicU32>,
71    #[allow(clippy::type_complexity)]
72    #[serde(skip)]
73    response_channels: Arc<
74        Mutex<
75            Option<
76                HashMap<
77                    u32,
78                    oneshot::Sender<(proto::Envelope, std::time::Instant, oneshot::Sender<()>)>,
79                >,
80            >,
81        >,
82    >,
83    #[allow(clippy::type_complexity)]
84    #[serde(skip)]
85    stream_response_channels: Arc<
86        Mutex<
87            Option<
88                HashMap<u32, mpsc::UnboundedSender<(Result<proto::Envelope>, oneshot::Sender<()>)>>,
89            >,
90        >,
91    >,
92}
93
94const KEEPALIVE_INTERVAL: Duration = Duration::from_secs(1);
95const WRITE_TIMEOUT: Duration = Duration::from_secs(2);
96pub const RECEIVE_TIMEOUT: Duration = Duration::from_secs(10);
97
98impl Peer {
99    pub fn new(epoch: u32) -> Arc<Self> {
100        Arc::new(Self {
101            epoch: AtomicU32::new(epoch),
102            connections: Default::default(),
103            next_connection_id: Default::default(),
104        })
105    }
106
107    pub fn epoch(&self) -> u32 {
108        self.epoch.load(SeqCst)
109    }
110
111    pub fn add_connection<F, Fut, Out>(
112        self: &Arc<Self>,
113        connection: Connection,
114        create_timer: F,
115    ) -> (
116        ConnectionId,
117        impl Future<Output = anyhow::Result<()>> + Send + use<F, Fut, Out>,
118        BoxStream<'static, Box<dyn AnyTypedEnvelope>>,
119    )
120    where
121        F: Send + Fn(Duration) -> Fut,
122        Fut: Send + Future<Output = Out>,
123        Out: Send,
124    {
125        // For outgoing messages, use an unbounded channel so that application code
126        // can always send messages without yielding. For incoming messages, use a
127        // bounded channel so that other peers will receive backpressure if they send
128        // messages faster than this peer can process them.
129        #[cfg(any(test, feature = "test-support"))]
130        const INCOMING_BUFFER_SIZE: usize = 1;
131        #[cfg(not(any(test, feature = "test-support")))]
132        const INCOMING_BUFFER_SIZE: usize = 256;
133        let (mut incoming_tx, incoming_rx) = mpsc::channel(INCOMING_BUFFER_SIZE);
134        let (outgoing_tx, mut outgoing_rx) = mpsc::unbounded();
135
136        let connection_id = ConnectionId {
137            owner_id: self.epoch.load(SeqCst),
138            id: self.next_connection_id.fetch_add(1, SeqCst),
139        };
140        let connection_state = ConnectionState {
141            outgoing_tx,
142            next_message_id: Default::default(),
143            response_channels: Arc::new(Mutex::new(Some(Default::default()))),
144            stream_response_channels: Arc::new(Mutex::new(Some(Default::default()))),
145        };
146        let mut writer = MessageStream::new(connection.tx);
147        let mut reader = MessageStream::new(connection.rx);
148
149        let this = self.clone();
150        let response_channels = connection_state.response_channels.clone();
151        let stream_response_channels = connection_state.stream_response_channels.clone();
152
153        let handle_io = async move {
154            tracing::trace!(%connection_id, "handle io future: start");
155
156            let _end_connection = util::defer(|| {
157                response_channels.lock().take();
158                if let Some(channels) = stream_response_channels.lock().take() {
159                    for channel in channels.values() {
160                        let _ = channel.unbounded_send((
161                            Err(anyhow!("connection closed")),
162                            oneshot::channel().0,
163                        ));
164                    }
165                }
166                this.connections.write().remove(&connection_id);
167                tracing::trace!(%connection_id, "handle io future: end");
168            });
169
170            // Send messages on this frequency so the connection isn't closed.
171            let keepalive_timer = create_timer(KEEPALIVE_INTERVAL).fuse();
172            futures::pin_mut!(keepalive_timer);
173
174            // Disconnect if we don't receive messages at least this frequently.
175            let receive_timeout = create_timer(RECEIVE_TIMEOUT).fuse();
176            futures::pin_mut!(receive_timeout);
177
178            loop {
179                tracing::trace!(%connection_id, "outer loop iteration start");
180                let read_message = reader.read().fuse();
181                futures::pin_mut!(read_message);
182
183                loop {
184                    tracing::trace!(%connection_id, "inner loop iteration start");
185                    futures::select_biased! {
186                        outgoing = outgoing_rx.next().fuse() => match outgoing {
187                            Some(outgoing) => {
188                                tracing::trace!(%connection_id, "outgoing rpc message: writing");
189                                futures::select_biased! {
190                                    result = writer.write(outgoing).fuse() => {
191                                        tracing::trace!(%connection_id, "outgoing rpc message: done writing");
192                                        result.context("failed to write RPC message")?;
193                                        tracing::trace!(%connection_id, "keepalive interval: resetting after sending message");
194                                        keepalive_timer.set(create_timer(KEEPALIVE_INTERVAL).fuse());
195                                    }
196                                    _ = create_timer(WRITE_TIMEOUT).fuse() => {
197                                        tracing::trace!(%connection_id, "outgoing rpc message: writing timed out");
198                                        anyhow::bail!("timed out writing message");
199                                    }
200                                }
201                            }
202                            None => {
203                                tracing::trace!(%connection_id, "outgoing rpc message: channel closed");
204                                return Ok(())
205                            },
206                        },
207                        _ = keepalive_timer => {
208                            tracing::trace!(%connection_id, "keepalive interval: pinging");
209                            futures::select_biased! {
210                                result = writer.write(Message::Ping).fuse() => {
211                                    tracing::trace!(%connection_id, "keepalive interval: done pinging");
212                                    result.context("failed to send keepalive")?;
213                                    tracing::trace!(%connection_id, "keepalive interval: resetting after pinging");
214                                    keepalive_timer.set(create_timer(KEEPALIVE_INTERVAL).fuse());
215                                }
216                                _ = create_timer(WRITE_TIMEOUT).fuse() => {
217                                    tracing::trace!(%connection_id, "keepalive interval: pinging timed out");
218                                    anyhow::bail!("timed out sending keepalive");
219                                }
220                            }
221                        }
222                        incoming = read_message => {
223                            let incoming = incoming.context("error reading rpc message from socket")?;
224                            tracing::trace!(%connection_id, "incoming rpc message: received");
225                            tracing::trace!(%connection_id, "receive timeout: resetting");
226                            receive_timeout.set(create_timer(RECEIVE_TIMEOUT).fuse());
227                            if let (Message::Envelope(incoming), received_at) = incoming {
228                                tracing::trace!(%connection_id, "incoming rpc message: processing");
229                                futures::select_biased! {
230                                    result = incoming_tx.send((incoming, received_at)).fuse() => match result {
231                                        Ok(_) => {
232                                            tracing::trace!(%connection_id, "incoming rpc message: processed");
233                                        }
234                                        Err(_) => {
235                                            tracing::trace!(%connection_id, "incoming rpc message: channel closed");
236                                            return Ok(())
237                                        }
238                                    },
239                                    _ = create_timer(WRITE_TIMEOUT).fuse() => {
240                                        tracing::trace!(%connection_id, "incoming rpc message: processing timed out");
241                                        anyhow::bail!("timed out processing incoming message");
242                                    }
243                                }
244                            }
245                            break;
246                        },
247                        _ = receive_timeout => {
248                            tracing::trace!(%connection_id, "receive timeout: delay between messages too long");
249                            anyhow::bail!("delay between messages too long");
250                        }
251                    }
252                }
253            }
254        };
255
256        let response_channels = connection_state.response_channels.clone();
257        let stream_response_channels = connection_state.stream_response_channels.clone();
258        self.connections
259            .write()
260            .insert(connection_id, connection_state);
261
262        let incoming_rx = incoming_rx.filter_map(move |(incoming, received_at)| {
263            let response_channels = response_channels.clone();
264            let stream_response_channels = stream_response_channels.clone();
265            async move {
266                let message_id = incoming.id;
267                tracing::trace!(?incoming, "incoming message future: start");
268                let _end = util::defer(move || {
269                    tracing::trace!(%connection_id, message_id, "incoming message future: end");
270                });
271
272                if let Some(responding_to) = incoming.responding_to {
273                    tracing::trace!(
274                        %connection_id,
275                        message_id,
276                        responding_to,
277                        "incoming response: received"
278                    );
279                    let response_channel =
280                        response_channels.lock().as_mut()?.remove(&responding_to);
281                    let terminal_stream_response = matches!(
282                        &incoming.payload,
283                        Some(proto::envelope::Payload::Error(_))
284                            | Some(proto::envelope::Payload::EndStream(_))
285                    );
286                    let stream_response_channel = if terminal_stream_response {
287                        stream_response_channels
288                            .lock()
289                            .as_mut()?
290                            .remove(&responding_to)
291                    } else {
292                        stream_response_channels
293                            .lock()
294                            .as_ref()?
295                            .get(&responding_to)
296                            .cloned()
297                    };
298
299                    if let Some(tx) = response_channel {
300                        let requester_resumed = oneshot::channel();
301                        if let Err(error) = tx.send((incoming, received_at, requester_resumed.0)) {
302                            tracing::trace!(
303                                %connection_id,
304                                message_id,
305                                responding_to = responding_to,
306                                ?error,
307                                "incoming response: request future dropped",
308                            );
309                        }
310
311                        tracing::trace!(
312                            %connection_id,
313                            message_id,
314                            responding_to,
315                            "incoming response: waiting to resume requester"
316                        );
317                        let _ = requester_resumed.1.await;
318                        tracing::trace!(
319                            %connection_id,
320                            message_id,
321                            responding_to,
322                            "incoming response: requester resumed"
323                        );
324                    } else if let Some(tx) = stream_response_channel {
325                        let requester_resumed = oneshot::channel();
326                        if let Err(error) = tx.unbounded_send((Ok(incoming), requester_resumed.0)) {
327                            tracing::debug!(
328                                %connection_id,
329                                message_id,
330                                responding_to = responding_to,
331                                ?error,
332                                "incoming stream response: request future dropped",
333                            );
334                            // The consumer has gone away, so drop the bookkeeping
335                            // for this stream rather than letting it accumulate
336                            // every subsequent message until a terminal frame.
337                            if let Some(channels) = stream_response_channels.lock().as_mut() {
338                                channels.remove(&responding_to);
339                            }
340                        } else {
341                            let _ = requester_resumed.1.await;
342                        }
343                    } else {
344                        let message_type = proto::build_typed_envelope(
345                            connection_id.into(),
346                            received_at,
347                            incoming,
348                        )
349                        .map(|p| p.payload_type_name());
350                        tracing::warn!(
351                            %connection_id,
352                            message_id,
353                            responding_to,
354                            message_type,
355                            "incoming response: unknown request"
356                        );
357                    }
358
359                    None
360                } else {
361                    tracing::trace!(%connection_id, message_id, "incoming message: received");
362                    proto::build_typed_envelope(connection_id.into(), received_at, incoming)
363                        .or_else(|| {
364                            tracing::error!(
365                                %connection_id,
366                                message_id,
367                                "unable to construct a typed envelope"
368                            );
369                            None
370                        })
371                }
372            }
373        });
374        (connection_id, handle_io, incoming_rx.boxed())
375    }
376
377    #[cfg(any(test, feature = "test-support"))]
378    pub fn add_test_connection(
379        self: &Arc<Self>,
380        connection: Connection,
381        executor: gpui::BackgroundExecutor,
382    ) -> (
383        ConnectionId,
384        impl Future<Output = anyhow::Result<()>> + Send + use<>,
385        BoxStream<'static, Box<dyn AnyTypedEnvelope>>,
386    ) {
387        self.add_connection(connection, move |duration| executor.timer(duration))
388    }
389
390    pub fn disconnect(&self, connection_id: ConnectionId) {
391        self.connections.write().remove(&connection_id);
392    }
393
394    #[cfg(any(test, feature = "test-support"))]
395    pub fn reset(&self, epoch: u32) {
396        self.next_connection_id.store(0, SeqCst);
397        self.epoch.store(epoch, SeqCst);
398    }
399
400    pub fn teardown(&self) {
401        self.connections.write().clear();
402    }
403
404    /// Make a request and wait for a response.
405    pub fn request<T: RequestMessage>(
406        &self,
407        receiver_id: ConnectionId,
408        request: T,
409    ) -> impl Future<Output = Result<T::Response>> + use<T> {
410        self.request_internal(None, receiver_id, request)
411            .map_ok(|envelope| envelope.payload)
412    }
413
414    pub fn request_envelope<T: RequestMessage>(
415        &self,
416        receiver_id: ConnectionId,
417        request: T,
418    ) -> impl Future<Output = Result<TypedEnvelope<T::Response>>> + use<T> {
419        self.request_internal(None, receiver_id, request)
420    }
421
422    pub fn forward_request<T: RequestMessage>(
423        &self,
424        sender_id: ConnectionId,
425        receiver_id: ConnectionId,
426        request: T,
427    ) -> impl Future<Output = Result<T::Response>> {
428        self.request_internal(Some(sender_id), receiver_id, request)
429            .map_ok(|envelope| envelope.payload)
430    }
431
432    fn request_internal<T: RequestMessage>(
433        &self,
434        original_sender_id: Option<ConnectionId>,
435        receiver_id: ConnectionId,
436        request: T,
437    ) -> impl Future<Output = Result<TypedEnvelope<T::Response>>> + use<T> {
438        let envelope = request.into_envelope(0, None, original_sender_id.map(Into::into));
439        let response = self.request_dynamic(receiver_id, envelope, T::NAME);
440        async move {
441            let (response, received_at) = response.await?;
442            Ok(TypedEnvelope {
443                message_id: response.id,
444                sender_id: receiver_id.into(),
445                original_sender_id: response.original_sender_id,
446                payload: T::Response::from_envelope(response)
447                    .context("received response of the wrong type")?,
448                received_at,
449            })
450        }
451    }
452
453    /// Make a request and wait for a response.
454    ///
455    /// The caller must make sure to deserialize the response into the request's
456    /// response type. This interface is only useful in trait objects, where
457    /// generics can't be used. If you have a concrete type, use `request`.
458    pub fn request_dynamic(
459        &self,
460        receiver_id: ConnectionId,
461        mut envelope: proto::Envelope,
462        type_name: &'static str,
463    ) -> impl Future<Output = Result<(proto::Envelope, Instant)>> + use<> {
464        let (tx, rx) = oneshot::channel();
465        let send = self.connection_state(receiver_id).and_then(|connection| {
466            envelope.id = connection.next_message_id.fetch_add(1, SeqCst);
467            connection
468                .response_channels
469                .lock()
470                .as_mut()
471                .context("connection was closed")?
472                // requesting to forward the response on the oneshot tx
473                // when it's envelope.id matches the one for the request we are about to send
474                .insert(envelope.id, tx);
475            connection
476                .outgoing_tx
477                // request that the message is send at some point in the future
478                .unbounded_send(Message::Envelope(envelope))
479                .context("connection was closed")?;
480            Ok(())
481        });
482        async move {
483            send?; // wait for reception
484            let (response, received_at, _barrier) = rx.await.context("connection was closed")?;
485            if let Some(proto::envelope::Payload::Error(error)) = &response.payload {
486                return Err(RpcError::from_proto(error, type_name));
487            }
488            Ok((response, received_at))
489        }
490    }
491
492    pub fn request_stream<T: RequestMessage>(
493        &self,
494        receiver_id: ConnectionId,
495        request: T,
496    ) -> impl Future<Output = Result<BoxStream<'static, Result<T::Response>>>> {
497        let stream =
498            self.request_stream_dynamic(receiver_id, request.into_envelope(0, None, None), T::NAME);
499
500        async move {
501            Ok(stream
502                .await?
503                .map(|response| {
504                    T::Response::from_envelope(response?)
505                        .context("received response of the wrong type")
506                })
507                .boxed())
508        }
509    }
510
511    pub fn request_stream_dynamic(
512        &self,
513        receiver_id: ConnectionId,
514        mut envelope: proto::Envelope,
515        request_type: &'static str,
516    ) -> impl Future<Output = Result<BoxStream<'static, Result<proto::Envelope>>>> + use<> {
517        let (tx, rx) = mpsc::unbounded();
518        let send = self.connection_state(receiver_id).and_then(|connection| {
519            let message_id = connection.next_message_id.fetch_add(1, SeqCst);
520            envelope.id = message_id;
521            let stream_response_channels = connection.stream_response_channels.clone();
522            stream_response_channels
523                .lock()
524                .as_mut()
525                .context("connection was closed")?
526                .insert(message_id, tx);
527            if let Err(error) = connection
528                .outgoing_tx
529                .unbounded_send(Message::Envelope(envelope))
530            {
531                if let Some(channels) = stream_response_channels.lock().as_mut() {
532                    channels.remove(&message_id);
533                }
534                return Err(error).context("connection was closed");
535            }
536            Ok((message_id, stream_response_channels))
537        });
538
539        async move {
540            let (message_id, stream_response_channels) = send?;
541            let stream_response_channels = Arc::downgrade(&stream_response_channels);
542            let cleanup_stream_response_channel = util::defer({
543                let stream_response_channels = stream_response_channels.clone();
544                move || {
545                    if let Some(channels) = stream_response_channels.upgrade()
546                        && let Some(channels) = channels.lock().as_mut()
547                    {
548                        channels.remove(&message_id);
549                    }
550                }
551            });
552
553            Ok(rx
554                .filter_map(move |(response, _barrier)| {
555                    let _keep_cleanup_guard_alive = &cleanup_stream_response_channel;
556                    let stream_response_channels = stream_response_channels.clone();
557                    future::ready(match response {
558                        Ok(response) => {
559                            if let Some(proto::envelope::Payload::Error(error)) = &response.payload
560                            {
561                                // Remove the transmitting end of the response channel to end the stream.
562                                if let Some(channels) = stream_response_channels.upgrade()
563                                    && let Some(channels) = channels.lock().as_mut()
564                                {
565                                    channels.remove(&message_id);
566                                }
567                                Some(Err(RpcError::from_proto(error, request_type)))
568                            } else if let Some(proto::envelope::Payload::EndStream(_)) =
569                                &response.payload
570                            {
571                                // Remove the transmitting end of the response channel to end the stream.
572                                if let Some(channels) = stream_response_channels.upgrade()
573                                    && let Some(channels) = channels.lock().as_mut()
574                                {
575                                    channels.remove(&message_id);
576                                }
577                                None
578                            } else {
579                                Some(Ok(response))
580                            }
581                        }
582                        Err(error) => Some(Err(error)),
583                    })
584                })
585                .boxed())
586        }
587    }
588
589    pub fn send<T: EnvelopedMessage>(&self, receiver_id: ConnectionId, message: T) -> Result<()> {
590        let connection = self.connection_state(receiver_id)?;
591        let message_id = connection
592            .next_message_id
593            .fetch_add(1, atomic::Ordering::SeqCst);
594        connection.outgoing_tx.unbounded_send(Message::Envelope(
595            message.into_envelope(message_id, None, None),
596        ))?;
597        Ok(())
598    }
599
600    pub fn send_dynamic(&self, receiver_id: ConnectionId, message: proto::Envelope) -> Result<()> {
601        let connection = self.connection_state(receiver_id)?;
602        connection
603            .outgoing_tx
604            .unbounded_send(Message::Envelope(message))?;
605        Ok(())
606    }
607
608    pub fn forward_send<T: EnvelopedMessage>(
609        &self,
610        sender_id: ConnectionId,
611        receiver_id: ConnectionId,
612        message: T,
613    ) -> Result<()> {
614        let connection = self.connection_state(receiver_id)?;
615        let message_id = connection
616            .next_message_id
617            .fetch_add(1, atomic::Ordering::SeqCst);
618        connection
619            .outgoing_tx
620            .unbounded_send(Message::Envelope(message.into_envelope(
621                message_id,
622                None,
623                Some(sender_id.into()),
624            )))?;
625        Ok(())
626    }
627
628    pub fn respond<T: RequestMessage>(
629        &self,
630        receipt: Receipt<T>,
631        response: T::Response,
632    ) -> Result<()> {
633        let connection = self.connection_state(receipt.sender_id.into())?;
634        let message_id = connection
635            .next_message_id
636            .fetch_add(1, atomic::Ordering::SeqCst);
637        connection
638            .outgoing_tx
639            .unbounded_send(Message::Envelope(response.into_envelope(
640                message_id,
641                Some(receipt.message_id),
642                None,
643            )))?;
644        Ok(())
645    }
646
647    pub fn end_stream<T: RequestMessage>(&self, receipt: Receipt<T>) -> Result<()> {
648        let connection = self.connection_state(receipt.sender_id.into())?;
649        let message_id = connection
650            .next_message_id
651            .fetch_add(1, atomic::Ordering::SeqCst);
652
653        let message = proto::EndStream {};
654
655        connection
656            .outgoing_tx
657            .unbounded_send(Message::Envelope(message.into_envelope(
658                message_id,
659                Some(receipt.message_id),
660                None,
661            )))?;
662        Ok(())
663    }
664
665    pub fn respond_with_error<T: RequestMessage>(
666        &self,
667        receipt: Receipt<T>,
668        response: proto::Error,
669    ) -> Result<()> {
670        let connection = self.connection_state(receipt.sender_id.into())?;
671        let message_id = connection
672            .next_message_id
673            .fetch_add(1, atomic::Ordering::SeqCst);
674        connection
675            .outgoing_tx
676            .unbounded_send(Message::Envelope(response.into_envelope(
677                message_id,
678                Some(receipt.message_id),
679                None,
680            )))?;
681        Ok(())
682    }
683
684    pub fn respond_with_unhandled_message(
685        &self,
686        sender_id: ConnectionId,
687        request_message_id: u32,
688        message_type_name: &'static str,
689    ) -> Result<()> {
690        let connection = self.connection_state(sender_id)?;
691        let response = ErrorCode::Internal
692            .message(format!("message {} was not handled", message_type_name))
693            .to_proto();
694        let message_id = connection
695            .next_message_id
696            .fetch_add(1, atomic::Ordering::SeqCst);
697        connection
698            .outgoing_tx
699            .unbounded_send(Message::Envelope(response.into_envelope(
700                message_id,
701                Some(request_message_id),
702                None,
703            )))?;
704        Ok(())
705    }
706
707    fn connection_state(&self, connection_id: ConnectionId) -> Result<ConnectionState> {
708        let connections = self.connections.read();
709        let connection = connections
710            .get(&connection_id)
711            .with_context(|| format!("no such connection: {connection_id}"))?;
712        Ok(connection.clone())
713    }
714
715    #[cfg(any(test, feature = "test-support"))]
716    pub fn pending_stream_request_count(&self, connection_id: ConnectionId) -> Option<usize> {
717        let connection = self.connection_state(connection_id).ok()?;
718        let channels = connection.stream_response_channels.lock();
719        Some(channels.as_ref()?.len())
720    }
721}
722
723impl Serialize for Peer {
724    fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
725    where
726        S: serde::Serializer,
727    {
728        let mut state = serializer.serialize_struct("Peer", 2)?;
729        state.serialize_field("connections", &*self.connections.read())?;
730        state.end()
731    }
732}
733
734#[cfg(test)]
735mod tests {
736    use super::*;
737    use async_tungstenite::tungstenite::Message as WebSocketMessage;
738    use gpui::TestAppContext;
739
740    fn init_logger() {
741        zlog::init_test();
742    }
743
744    #[gpui::test(iterations = 50)]
745    async fn test_request_response(cx: &mut TestAppContext) {
746        init_logger();
747
748        let executor = cx.executor();
749
750        // create 2 clients connected to 1 server
751        let server = Peer::new(0);
752        let client1 = Peer::new(0);
753        let client2 = Peer::new(0);
754
755        let (client1_to_server_conn, server_to_client_1_conn, _kill) =
756            Connection::in_memory(cx.executor());
757        let (client1_conn_id, io_task1, client1_incoming) =
758            client1.add_test_connection(client1_to_server_conn, cx.executor());
759        let (_, io_task2, server_incoming1) =
760            server.add_test_connection(server_to_client_1_conn, cx.executor());
761
762        let (client2_to_server_conn, server_to_client_2_conn, _kill) =
763            Connection::in_memory(cx.executor());
764        let (client2_conn_id, io_task3, client2_incoming) =
765            client2.add_test_connection(client2_to_server_conn, cx.executor());
766        let (_, io_task4, server_incoming2) =
767            server.add_test_connection(server_to_client_2_conn, cx.executor());
768
769        executor.spawn(io_task1).detach();
770        executor.spawn(io_task2).detach();
771        executor.spawn(io_task3).detach();
772        executor.spawn(io_task4).detach();
773        executor
774            .spawn(handle_messages(server_incoming1, server.clone()))
775            .detach();
776        executor
777            .spawn(handle_messages(client1_incoming, client1.clone()))
778            .detach();
779        executor
780            .spawn(handle_messages(server_incoming2, server.clone()))
781            .detach();
782        executor
783            .spawn(handle_messages(client2_incoming, client2.clone()))
784            .detach();
785
786        assert_eq!(
787            client1
788                .request(client1_conn_id, proto::Ping {},)
789                .await
790                .unwrap(),
791            proto::Ack {}
792        );
793
794        assert_eq!(
795            client2
796                .request(client2_conn_id, proto::Ping {},)
797                .await
798                .unwrap(),
799            proto::Ack {}
800        );
801
802        assert_eq!(
803            client1
804                .request(client1_conn_id, proto::Test { id: 1 },)
805                .await
806                .unwrap(),
807            proto::Test { id: 1 }
808        );
809
810        assert_eq!(
811            client2
812                .request(client2_conn_id, proto::Test { id: 2 })
813                .await
814                .unwrap(),
815            proto::Test { id: 2 }
816        );
817
818        client1.disconnect(client1_conn_id);
819        client2.disconnect(client1_conn_id);
820
821        async fn handle_messages(
822            mut messages: BoxStream<'static, Box<dyn AnyTypedEnvelope>>,
823            peer: Arc<Peer>,
824        ) -> Result<()> {
825            while let Some(envelope) = messages.next().await {
826                let envelope = envelope.into_any();
827                if let Some(envelope) = envelope.downcast_ref::<TypedEnvelope<proto::Ping>>() {
828                    let receipt = envelope.receipt();
829                    peer.respond(receipt, proto::Ack {})?
830                } else if let Some(envelope) = envelope.downcast_ref::<TypedEnvelope<proto::Test>>()
831                {
832                    peer.respond(envelope.receipt(), envelope.payload.clone())?
833                } else {
834                    panic!("unknown message type");
835                }
836            }
837
838            Ok(())
839        }
840    }
841
842    #[gpui::test(iterations = 50)]
843    async fn test_order_of_response_and_incoming(cx: &mut TestAppContext) {
844        let executor = cx.executor();
845        let server = Peer::new(0);
846        let client = Peer::new(0);
847
848        let (client_to_server_conn, server_to_client_conn, _kill) =
849            Connection::in_memory(executor.clone());
850        let (client_to_server_conn_id, io_task1, mut client_incoming) =
851            client.add_test_connection(client_to_server_conn, executor.clone());
852
853        let (server_to_client_conn_id, io_task2, mut server_incoming) =
854            server.add_test_connection(server_to_client_conn, executor.clone());
855
856        executor.spawn(io_task1).detach();
857        executor.spawn(io_task2).detach();
858
859        executor
860            .spawn(async move {
861                let future = server_incoming.next().await;
862                let request = future
863                    .unwrap()
864                    .into_any()
865                    .downcast::<TypedEnvelope<proto::Ping>>()
866                    .unwrap();
867
868                server
869                    .send(
870                        server_to_client_conn_id,
871                        ErrorCode::Internal
872                            .message("message 1".to_string())
873                            .to_proto(),
874                    )
875                    .unwrap();
876                server
877                    .send(
878                        server_to_client_conn_id,
879                        ErrorCode::Internal
880                            .message("message 2".to_string())
881                            .to_proto(),
882                    )
883                    .unwrap();
884                server.respond(request.receipt(), proto::Ack {}).unwrap();
885
886                // Prevent the connection from being dropped
887                server_incoming.next().await;
888            })
889            .detach();
890
891        let events = Arc::new(Mutex::new(Vec::new()));
892
893        let response = client.request(client_to_server_conn_id, proto::Ping {});
894        let response_task = executor.spawn({
895            let events = events.clone();
896            async move {
897                response.await.unwrap();
898                events.lock().push("response".to_string());
899            }
900        });
901
902        executor
903            .spawn({
904                let events = events.clone();
905                async move {
906                    let incoming1 = client_incoming
907                        .next()
908                        .await
909                        .unwrap()
910                        .into_any()
911                        .downcast::<TypedEnvelope<proto::Error>>()
912                        .unwrap();
913                    events.lock().push(incoming1.payload.message);
914                    let incoming2 = client_incoming
915                        .next()
916                        .await
917                        .unwrap()
918                        .into_any()
919                        .downcast::<TypedEnvelope<proto::Error>>()
920                        .unwrap();
921                    events.lock().push(incoming2.payload.message);
922
923                    // Prevent the connection from being dropped
924                    client_incoming.next().await;
925                }
926            })
927            .detach();
928
929        response_task.await;
930        assert_eq!(
931            &*events.lock(),
932            &[
933                "message 1".to_string(),
934                "message 2".to_string(),
935                "response".to_string()
936            ]
937        );
938    }
939
940    #[gpui::test(iterations = 50)]
941    async fn test_dropping_request_before_completion(cx: &mut TestAppContext) {
942        let executor = cx.executor();
943        let server = Peer::new(0);
944        let client = Peer::new(0);
945
946        let (client_to_server_conn, server_to_client_conn, _kill) =
947            Connection::in_memory(cx.executor());
948        let (client_to_server_conn_id, io_task1, mut client_incoming) =
949            client.add_test_connection(client_to_server_conn, cx.executor());
950        let (server_to_client_conn_id, io_task2, mut server_incoming) =
951            server.add_test_connection(server_to_client_conn, cx.executor());
952
953        executor.spawn(io_task1).detach();
954        executor.spawn(io_task2).detach();
955
956        executor
957            .spawn(async move {
958                let request1 = server_incoming
959                    .next()
960                    .await
961                    .unwrap()
962                    .into_any()
963                    .downcast::<TypedEnvelope<proto::Ping>>()
964                    .unwrap();
965                let request2 = server_incoming
966                    .next()
967                    .await
968                    .unwrap()
969                    .into_any()
970                    .downcast::<TypedEnvelope<proto::Ping>>()
971                    .unwrap();
972
973                server
974                    .send(
975                        server_to_client_conn_id,
976                        ErrorCode::Internal
977                            .message("message 1".to_string())
978                            .to_proto(),
979                    )
980                    .unwrap();
981                server
982                    .send(
983                        server_to_client_conn_id,
984                        ErrorCode::Internal
985                            .message("message 2".to_string())
986                            .to_proto(),
987                    )
988                    .unwrap();
989                server.respond(request1.receipt(), proto::Ack {}).unwrap();
990                server.respond(request2.receipt(), proto::Ack {}).unwrap();
991
992                // Prevent the connection from being dropped
993                server_incoming.next().await;
994            })
995            .detach();
996
997        let events = Arc::new(Mutex::new(Vec::new()));
998
999        let request1 = client.request(client_to_server_conn_id, proto::Ping {});
1000        let request1_task = executor.spawn(request1);
1001        let request2 = client.request(client_to_server_conn_id, proto::Ping {});
1002        let request2_task = executor.spawn({
1003            let events = events.clone();
1004            async move {
1005                request2.await.unwrap();
1006                events.lock().push("response 2".to_string());
1007            }
1008        });
1009
1010        executor
1011            .spawn({
1012                let events = events.clone();
1013                async move {
1014                    let incoming1 = client_incoming
1015                        .next()
1016                        .await
1017                        .unwrap()
1018                        .into_any()
1019                        .downcast::<TypedEnvelope<proto::Error>>()
1020                        .unwrap();
1021                    events.lock().push(incoming1.payload.message);
1022                    let incoming2 = client_incoming
1023                        .next()
1024                        .await
1025                        .unwrap()
1026                        .into_any()
1027                        .downcast::<TypedEnvelope<proto::Error>>()
1028                        .unwrap();
1029                    events.lock().push(incoming2.payload.message);
1030
1031                    // Prevent the connection from being dropped
1032                    client_incoming.next().await;
1033                }
1034            })
1035            .detach();
1036
1037        // Allow the request to make some progress before dropping it.
1038        cx.executor().simulate_random_delay().await;
1039        drop(request1_task);
1040
1041        request2_task.await;
1042        assert_eq!(
1043            &*events.lock(),
1044            &[
1045                "message 1".to_string(),
1046                "message 2".to_string(),
1047                "response 2".to_string()
1048            ]
1049        );
1050    }
1051
1052    #[gpui::test(iterations = 50)]
1053    async fn test_request_stream(cx: &mut TestAppContext) {
1054        init_logger();
1055
1056        let executor = cx.executor();
1057        let server = Peer::new(0);
1058        let client = Peer::new(0);
1059
1060        let (client_to_server_conn, server_to_client_conn, _kill) =
1061            Connection::in_memory(executor.clone());
1062        let (client_to_server_conn_id, io_task1, mut client_incoming) =
1063            client.add_test_connection(client_to_server_conn, executor.clone());
1064        let (_, io_task2, mut server_incoming) =
1065            server.add_test_connection(server_to_client_conn, executor.clone());
1066
1067        executor.spawn(io_task1).detach();
1068        executor.spawn(io_task2).detach();
1069        executor
1070            .spawn(async move { while client_incoming.next().await.is_some() {} })
1071            .detach();
1072
1073        executor
1074            .spawn({
1075                let server = server.clone();
1076                async move {
1077                    let request = server_incoming
1078                        .next()
1079                        .await
1080                        .unwrap()
1081                        .into_any()
1082                        .downcast::<TypedEnvelope<proto::Test>>()
1083                        .unwrap();
1084                    let receipt = request.receipt();
1085                    server.respond(receipt, proto::Test { id: 1 }).unwrap();
1086                    server.respond(receipt, proto::Test { id: 2 }).unwrap();
1087                    server.respond(receipt, proto::Test { id: 3 }).unwrap();
1088                    server.end_stream(receipt).unwrap();
1089
1090                    // Prevent the connection from being dropped.
1091                    server_incoming.next().await;
1092                }
1093            })
1094            .detach();
1095
1096        let mut stream = client
1097            .request_stream(client_to_server_conn_id, proto::Test { id: 0 })
1098            .await
1099            .unwrap();
1100
1101        let mut received = Vec::new();
1102        while let Some(item) = stream.next().await {
1103            received.push(item.unwrap());
1104        }
1105
1106        assert_eq!(
1107            received,
1108            vec![
1109                proto::Test { id: 1 },
1110                proto::Test { id: 2 },
1111                proto::Test { id: 3 },
1112            ]
1113        );
1114        assert_eq!(
1115            client.pending_stream_request_count(client_to_server_conn_id),
1116            Some(0)
1117        );
1118    }
1119
1120    #[gpui::test]
1121    async fn test_request_stream_send_failure_cleans_up_response_channel(cx: &mut TestAppContext) {
1122        init_logger();
1123
1124        let executor = cx.executor();
1125        let client = Peer::new(0);
1126
1127        let (client_to_server_conn, _server_to_client_conn, _kill) =
1128            Connection::in_memory(executor.clone());
1129        let (client_to_server_conn_id, io_task, _client_incoming) =
1130            client.add_test_connection(client_to_server_conn, executor.clone());
1131
1132        drop(io_task);
1133
1134        let result = client
1135            .request_stream(client_to_server_conn_id, proto::Test { id: 0 })
1136            .await;
1137
1138        assert!(
1139            result.is_err(),
1140            "stream request should fail when the connection write task has gone away"
1141        );
1142        assert_eq!(
1143            client.pending_stream_request_count(client_to_server_conn_id),
1144            Some(0),
1145            "failed stream request should not leave response channel bookkeeping behind"
1146        );
1147    }
1148
1149    #[gpui::test(iterations = 50)]
1150    async fn test_request_stream_terminates_on_error(cx: &mut TestAppContext) {
1151        init_logger();
1152
1153        let executor = cx.executor();
1154        let server = Peer::new(0);
1155        let client = Peer::new(0);
1156
1157        let (client_to_server_conn, server_to_client_conn, _kill) =
1158            Connection::in_memory(executor.clone());
1159        let (client_to_server_conn_id, io_task1, mut client_incoming) =
1160            client.add_test_connection(client_to_server_conn, executor.clone());
1161        let (_, io_task2, mut server_incoming) =
1162            server.add_test_connection(server_to_client_conn, executor.clone());
1163
1164        executor.spawn(io_task1).detach();
1165        executor.spawn(io_task2).detach();
1166        executor
1167            .spawn(async move { while client_incoming.next().await.is_some() {} })
1168            .detach();
1169
1170        executor
1171            .spawn({
1172                let server = server.clone();
1173                async move {
1174                    let request = server_incoming
1175                        .next()
1176                        .await
1177                        .unwrap()
1178                        .into_any()
1179                        .downcast::<TypedEnvelope<proto::Test>>()
1180                        .unwrap();
1181                    let receipt = request.receipt();
1182                    server.respond(receipt, proto::Test { id: 1 }).unwrap();
1183                    // Send an Error without a trailing EndStream. The Error alone
1184                    // should be treated as a terminal stream response.
1185                    server
1186                        .respond_with_error(
1187                            receipt,
1188                            ErrorCode::Internal.message("boom".to_string()).to_proto(),
1189                        )
1190                        .unwrap();
1191
1192                    // Prevent the connection from being dropped.
1193                    server_incoming.next().await;
1194                }
1195            })
1196            .detach();
1197
1198        let mut stream = client
1199            .request_stream(client_to_server_conn_id, proto::Test { id: 0 })
1200            .await
1201            .unwrap();
1202
1203        assert_eq!(stream.next().await.unwrap().unwrap(), proto::Test { id: 1 });
1204
1205        let error = stream.next().await.unwrap().unwrap_err();
1206        assert!(
1207            format!("{error}").contains("boom"),
1208            "expected error to surface server message, got: {error}"
1209        );
1210
1211        // The error alone (without an EndStream) should terminate the stream.
1212        assert!(stream.next().await.is_none());
1213        assert_eq!(
1214            client.pending_stream_request_count(client_to_server_conn_id),
1215            Some(0)
1216        );
1217    }
1218
1219    #[gpui::test(iterations = 50)]
1220    async fn test_dropping_stream_request_before_completion(cx: &mut TestAppContext) {
1221        init_logger();
1222
1223        let executor = cx.executor();
1224        let server = Peer::new(0);
1225        let client = Peer::new(0);
1226
1227        let (client_to_server_conn, server_to_client_conn, _kill) =
1228            Connection::in_memory(executor.clone());
1229        let (client_to_server_conn_id, io_task1, mut client_incoming) =
1230            client.add_test_connection(client_to_server_conn, executor.clone());
1231        let (_, io_task2, mut server_incoming) =
1232            server.add_test_connection(server_to_client_conn, executor.clone());
1233
1234        executor.spawn(io_task1).detach();
1235        executor.spawn(io_task2).detach();
1236        executor
1237            .spawn(async move { while client_incoming.next().await.is_some() {} })
1238            .detach();
1239
1240        let (drop_signal_tx, drop_signal_rx) = oneshot::channel::<()>();
1241        let server_task = executor.spawn({
1242            let server = server.clone();
1243            async move {
1244                let request = server_incoming
1245                    .next()
1246                    .await
1247                    .unwrap()
1248                    .into_any()
1249                    .downcast::<TypedEnvelope<proto::Test>>()
1250                    .unwrap();
1251                let receipt = request.receipt();
1252                server.respond(receipt, proto::Test { id: 1 }).unwrap();
1253
1254                // Wait until the consumer has dropped the stream.
1255                drop_signal_rx.await.ok();
1256
1257                // Send a non-terminal response after the consumer is gone. The
1258                // peer should detect that the receiver has been dropped and clean
1259                // up its bookkeeping. Crucially, we do NOT send EndStream here
1260                // because that would clean up via the terminal-response path and
1261                // mask the bug.
1262                server.respond(receipt, proto::Test { id: 2 }).unwrap();
1263
1264                // A Ping/Ack round-trip after the response acts as a sync
1265                // barrier: because messages over the in-memory connection are
1266                // delivered in order, by the time the client observes the Ack,
1267                // it has already processed the dropped response above.
1268                let ping = server_incoming
1269                    .next()
1270                    .await
1271                    .unwrap()
1272                    .into_any()
1273                    .downcast::<TypedEnvelope<proto::Ping>>()
1274                    .unwrap();
1275                server.respond(ping.receipt(), proto::Ack {}).unwrap();
1276
1277                // Prevent the connection from being dropped.
1278                server_incoming.next().await;
1279            }
1280        });
1281
1282        let mut stream = client
1283            .request_stream(client_to_server_conn_id, proto::Test { id: 0 })
1284            .await
1285            .unwrap();
1286
1287        assert_eq!(stream.next().await.unwrap().unwrap(), proto::Test { id: 1 });
1288
1289        // The stream is mid-flight, so the channel should be tracked.
1290        assert_eq!(
1291            client.pending_stream_request_count(client_to_server_conn_id),
1292            Some(1)
1293        );
1294
1295        drop(stream);
1296        drop_signal_tx.send(()).ok();
1297
1298        // Synchronization barrier: once this Ack arrives, the read loop has
1299        // already processed the orphaned stream response that came before it.
1300        client
1301            .request(client_to_server_conn_id, proto::Ping {})
1302            .await
1303            .unwrap();
1304
1305        assert_eq!(
1306            client.pending_stream_request_count(client_to_server_conn_id),
1307            Some(0),
1308            "stream channel should be removed once the consumer has dropped the stream"
1309        );
1310
1311        drop(server_task);
1312    }
1313
1314    #[gpui::test(iterations = 50)]
1315    async fn test_disconnect(cx: &mut TestAppContext) {
1316        let executor = cx.executor();
1317
1318        let (client_conn, mut server_conn, _kill) = Connection::in_memory(executor.clone());
1319
1320        let client = Peer::new(0);
1321        let (connection_id, io_handler, mut incoming) =
1322            client.add_test_connection(client_conn, executor.clone());
1323
1324        let (io_ended_tx, io_ended_rx) = oneshot::channel();
1325        executor
1326            .spawn(async move {
1327                io_handler.await.ok();
1328                io_ended_tx.send(()).unwrap();
1329            })
1330            .detach();
1331
1332        let (messages_ended_tx, messages_ended_rx) = oneshot::channel();
1333        executor
1334            .spawn(async move {
1335                incoming.next().await;
1336                messages_ended_tx.send(()).unwrap();
1337            })
1338            .detach();
1339
1340        client.disconnect(connection_id);
1341
1342        let _ = io_ended_rx.await;
1343        let _ = messages_ended_rx.await;
1344        assert!(
1345            server_conn
1346                .send(WebSocketMessage::Binary(vec![].into()))
1347                .await
1348                .is_err()
1349        );
1350    }
1351
1352    #[gpui::test(iterations = 50)]
1353    async fn test_io_error(cx: &mut TestAppContext) {
1354        let executor = cx.executor();
1355        let (client_conn, mut server_conn, _kill) = Connection::in_memory(executor.clone());
1356
1357        let client = Peer::new(0);
1358        let (connection_id, io_handler, mut incoming) =
1359            client.add_test_connection(client_conn, executor.clone());
1360        executor.spawn(io_handler).detach();
1361        executor
1362            .spawn(async move { incoming.next().await })
1363            .detach();
1364
1365        let response = executor.spawn(client.request(connection_id, proto::Ping {}));
1366        let _request = server_conn.rx.next().await.unwrap().unwrap();
1367
1368        drop(server_conn);
1369        assert_eq!(
1370            response.await.unwrap_err().to_string(),
1371            "connection was closed"
1372        );
1373    }
1374}
1375