1use super::envelope::{Header, IncomingEnvelope, MessageKind, Parity, Side};
10use super::operation::{
11 OperationHandle, OperationKey, OutgoingBody, OutgoingMessage, PendingOperation,
12};
13use super::promise::{Notifications, PromiseResult};
14use super::worker;
15use super::{
16 Closer, DEFAULT_AUTOREPLY_TIMEOUT, DEFAULT_MAX_INBOUND_BYTES, DEFAULT_MAX_INBOUND_REQUESTS,
17 Error, Message, Promise, Requester, Responder, schema,
18};
19use crate::LogId;
20use crate::transport::{self, Read, Stream, Verifier, Write};
21use darkbio_clock::Clock;
22use darkbio_clock::sync::{Condvar, Mutex};
23use prost::bytes::Bytes;
24use std::collections::{HashMap, HashSet, VecDeque};
25use std::fmt;
26use std::sync::Arc;
27use std::sync::atomic::{AtomicUsize, Ordering};
28use std::time::{Duration, Instant};
29
30pub fn connect<R, W, V>(stream: Stream<R, W>, verifier: &V) -> Result<(Session, V::Info), Error>
36where
37 R: Read + Send + 'static,
38 W: Write + Send + 'static,
39 V: Verifier,
40{
41 let clock = stream.clock();
43 let mut client = transport::Client::new(stream);
44
45 let (sender, info) = client.connect(verifier)?;
47 #[cfg(any(test, feature = "fuzz"))]
48 let workers = Arc::new(worker::Tracker::default());
49 let session = Session::start(
50 Side::Client,
51 clock,
52 sender,
53 Some(client.closer()),
54 #[cfg(any(test, feature = "fuzz"))]
55 workers.clone(),
56 );
57 let inner = session.inner.clone();
58
59 worker::spawn(
61 "wire-client-reader",
62 #[cfg(any(test, feature = "fuzz"))]
63 &workers,
64 move || run_reader(client, inner),
65 );
66 Ok((session, info))
67}
68
69fn run_reader<R: Read, W: Write>(mut client: transport::Client<R, W>, session: Arc<SessionInner>) {
73 loop {
74 let result = client
75 .recv()
76 .map_err(Error::from)
77 .and_then(|bytes| session.handle_message(bytes));
78 if let Err(error) = result {
79 session.close(error);
82 break;
83 }
84 }
85}
86
87pub struct Session {
107 pub(super) inner: Arc<SessionInner>,
109}
110
111impl Session {
112 pub fn clock(&self) -> Clock {
115 self.inner.clock.clone()
116 }
117
118 pub fn set_autoreply_timeout(self, timeout: Duration) -> Self {
130 self.inner.set_autoreply_timeout(timeout);
131 self
132 }
133
134 pub fn set_inbound_limits(self, requests: usize, bytes: usize) -> Self {
153 let mut notifications = Notifications::default();
154 self.inner
155 .set_inbound_limits(requests, bytes, &mut notifications);
156 self
157 }
158
159 pub fn requester(&self) -> Requester {
161 Requester::new(Arc::downgrade(&self.inner), &self.inner.clock)
162 }
163
164 pub fn recv(&mut self) -> Result<(Message, Responder), Error> {
177 self.inner.recv()
178 }
179
180 pub fn closer(&self) -> Closer {
183 Closer::session(Arc::downgrade(&self.inner))
184 }
185
186 pub fn close(&self) {
192 self.inner.close(Error::Closed);
193 }
194}
195
196impl Drop for Session {
197 fn drop(&mut self) {
199 self.close();
200 }
201}
202
203impl fmt::Debug for Session {
204 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
207 let mut session = f.debug_struct("Session");
208 session.field("id", &self.inner.log_id);
209 if let Ok(state) = self.inner.state.try_lock() {
210 session.field("open", &matches!(*state, State::Open { .. }));
211 }
212 session.finish_non_exhaustive()
213 }
214}
215
216pub(super) struct SessionInner {
219 pub(super) clock: Clock,
221 state: Mutex<State>,
223 changed: Condvar,
225 retained_bytes: Arc<AtomicUsize>,
229 side: Side,
231 pub(super) log_id: LogId,
233 stream_closer: Option<transport::Closer>,
236 #[cfg(any(test, feature = "fuzz"))]
238 pub(super) workers: Arc<worker::Tracker>,
239 #[cfg(any(test, feature = "fuzz"))]
241 drop_hook: std::sync::Mutex<Option<std::sync::mpsc::Sender<()>>>,
242 #[cfg(any(test, feature = "fuzz"))]
244 disconnect_hook:
245 std::sync::Mutex<Option<(std::sync::mpsc::Sender<()>, std::sync::mpsc::Receiver<()>)>>,
246}
247
248#[cfg_attr(any(test, feature = "fuzz"), allow(clippy::large_enum_variant))]
252enum State {
253 Open {
256 max_inbound_requests: usize,
258 max_inbound_bytes: usize,
260 autoreply_timeout: Duration,
262
263 incoming: VecDeque<(u64, IncomingEnvelope)>,
265 reserved_ids: HashSet<u64>,
271
272 outgoing: VecDeque<OutgoingMessage>,
274 next_id: Option<u64>,
276 outstanding: HashMap<u64, OperationKey>,
279 operations: HashMap<OperationKey, PendingOperation>,
282
283 #[cfg(any(test, feature = "fuzz"))]
285 wait_hook: Option<std::sync::mpsc::Sender<()>>,
286 },
287 Closed(Error),
289}
290
291impl SessionInner {
292 fn new(
294 side: Side,
295 clock: Clock,
296 log_id: LogId,
297 stream_closer: Option<transport::Closer>,
298 #[cfg(any(test, feature = "fuzz"))] workers: Arc<worker::Tracker>,
299 ) -> Self {
300 Self {
301 changed: Condvar::new(&clock),
302 clock,
303 state: Mutex::new(State::Open {
304 max_inbound_requests: DEFAULT_MAX_INBOUND_REQUESTS,
305 max_inbound_bytes: DEFAULT_MAX_INBOUND_BYTES,
306 autoreply_timeout: DEFAULT_AUTOREPLY_TIMEOUT,
307
308 incoming: VecDeque::new(),
309 reserved_ids: HashSet::new(),
310
311 outgoing: VecDeque::new(),
312 next_id: Some(Parity::from(side).first()),
313 outstanding: HashMap::new(),
314 operations: HashMap::new(),
315
316 #[cfg(any(test, feature = "fuzz"))]
317 wait_hook: None,
318 }),
319 retained_bytes: Arc::new(AtomicUsize::new(0)),
320 side,
321 log_id,
322 stream_closer,
323 #[cfg(any(test, feature = "fuzz"))]
324 workers,
325 #[cfg(any(test, feature = "fuzz"))]
326 drop_hook: std::sync::Mutex::new(None),
327 #[cfg(any(test, feature = "fuzz"))]
328 disconnect_hook: std::sync::Mutex::new(None),
329 }
330 }
331
332 pub(super) fn set_autoreply_timeout(&self, timeout: Duration) {
335 let mut state = self.state.lock().expect("session state not poisoned");
336 if let State::Open {
337 autoreply_timeout, ..
338 } = &mut *state
339 {
340 *autoreply_timeout = timeout;
341 }
342 }
343
344 pub(super) fn set_inbound_limits(
347 &self,
348 requests: usize,
349 bytes: usize,
350 notifications: &mut Notifications,
351 ) {
352 let removed = {
354 let mut state = self.state.lock().expect("session state not poisoned");
355 let State::Open {
356 max_inbound_requests,
357 max_inbound_bytes,
358 reserved_ids,
359 ..
360 } = &mut *state
361 else {
362 return;
363 };
364 *max_inbound_requests = requests;
365 *max_inbound_bytes = bytes;
366 let error = if reserved_ids.len() > *max_inbound_requests {
367 Some(Error::InboundRequestLimitExceeded(*max_inbound_requests))
368 } else if self.retained_bytes.load(Ordering::Relaxed) > *max_inbound_bytes {
369 Some(Error::InboundByteLimitExceeded(*max_inbound_bytes))
370 } else {
371 None
372 };
373 error.and_then(|error| {
374 let (pending, queued) = state.workload();
375 let removed = state.close(error.clone(), self.now(), notifications);
376 if removed.is_some() {
377 self.log_closed(&error, pending, queued);
378 }
379 removed
380 })
381 };
382
383 if removed.is_some() {
385 self.finish_close(removed);
386 }
387 }
388
389 fn retain_incoming(
392 self: &Arc<Self>,
393 bytes: Bytes,
394 header: Header,
395 limit: usize,
396 ) -> Result<IncomingEnvelope, Error> {
397 IncomingEnvelope::new(
398 bytes,
399 header,
400 &self.retained_bytes,
401 limit,
402 self.side,
403 Arc::downgrade(self),
404 )
405 }
406
407 fn recv(self: &Arc<Self>) -> Result<(Message, Responder), Error> {
412 let mut state = self.state.lock().expect("session state not poisoned");
413 let (message, responder) = loop {
414 match &mut *state {
415 State::Closed(error) => return Err(error.clone()),
416 State::Open {
417 incoming,
418 #[cfg(any(test, feature = "fuzz"))]
419 wait_hook,
420 ..
421 } => {
422 if let Some((id, message)) = incoming.pop_front() {
423 break (
424 message,
425 Responder::new(Arc::downgrade(self), &self.clock, id),
426 );
427 }
428 #[cfg(any(test, feature = "fuzz"))]
429 if let Some(wait_hook) = wait_hook.take() {
430 let _ = wait_hook.send(());
431 }
432 state = self
433 .changed
434 .wait(state)
435 .expect("session state not poisoned");
436 }
437 }
438 };
439 drop(state);
440 Ok((message.decode()?, responder))
441 }
442
443 pub(super) fn close(&self, error: Error) {
447 let mut notifications = Notifications::default();
449 let (removed, pending, queued) = {
450 let mut state = self.state.lock().expect("session state not poisoned");
451 let (pending, queued) = state.workload();
452 (
453 state.close(error.clone(), self.now(), &mut notifications),
454 pending,
455 queued,
456 )
457 };
458
459 if removed.is_some() {
461 self.log_closed(&error, pending, queued);
462 }
463 self.finish_close(removed);
464 }
465
466 fn log_closed(&self, error: &Error, pending: usize, queued: usize) {
469 match error {
470 Error::Closed => tracing::info!(
471 "wire session {} closed locally (pending {}, queued {})",
472 self.log_id,
473 pending,
474 queued
475 ),
476 _ if error.orderly() => tracing::info!(
477 "wire session {} closed: {} (pending {}, queued {})",
478 self.log_id,
479 error.reason(),
480 pending,
481 queued
482 ),
483 _ => tracing::warn!(
484 "wire session {} failed: {} (pending {}, queued {})",
485 self.log_id,
486 error.reason(),
487 pending,
488 queued
489 ),
490 }
491 }
492
493 fn finish_close(&self, removed: Option<State>) {
495 self.changed.notify_all();
498 drop(removed);
499 if let Some(stream_closer) = &self.stream_closer {
500 stream_closer.close();
501 }
502 }
503
504 pub(super) fn reply_unanswered(self: &Arc<Self>, id: u64) {
507 self.autoreply(
508 id,
509 "unanswered",
510 schema::Error::reserved(
511 schema::ReservedErrors::Unanswered,
512 "request left unanswered",
513 ),
514 );
515 }
516
517 fn reply_unknown(self: &Arc<Self>, id: u64) {
520 self.autoreply(
521 id,
522 "unknown",
523 schema::Error::reserved(schema::ReservedErrors::Unknown, "request not known"),
524 );
525 }
526
527 fn autoreply(self: &Arc<Self>, id: u64, reason: &str, error: schema::Error) {
532 let now = self.now();
533 let timeout = {
534 let state = self.state.lock().expect("session state not poisoned");
535 match &*state {
536 State::Open {
537 autoreply_timeout, ..
538 } => *autoreply_timeout,
539 State::Closed(_) => return,
540 }
541 };
542 tracing::debug!("answering request {} as {}", id, reason);
543
544 let deadline = now.checked_add(timeout).unwrap_or(now);
547 let _ = self.reply(id, Err(error), deadline);
548 }
549
550 pub(super) fn request(
553 self: &Arc<Self>,
554 request: Message,
555 deadline: Instant,
556 ) -> Result<Promise<Message>, Error> {
557 let (sender, promise) = Promise::pair(Arc::downgrade(self), deadline, true);
558 self.enqueue(
559 OutgoingBody::Request(request),
560 PendingOperation {
561 deadline,
562 sender,
563 log_id: None,
564 },
565 )?;
566 Ok(promise)
567 }
568
569 pub(super) fn reply(
572 self: &Arc<Self>,
573 id: u64,
574 result: Result<Message, schema::Error>,
575 deadline: Instant,
576 ) -> Result<Promise<()>, Error> {
577 let (sender, promise) = Promise::pair(Arc::downgrade(self), deadline, false);
578 self.enqueue(
579 OutgoingBody::Reply { id, result },
580 PendingOperation {
581 deadline,
582 sender,
583 log_id: Some(LogId::from(id)),
584 },
585 )?;
586 Ok(promise)
587 }
588
589 fn enqueue(
593 self: &Arc<Self>,
594 body: OutgoingBody,
595 operation: PendingOperation,
596 ) -> Result<(), Error> {
597 let mut notifications = Notifications::default();
599 {
600 let mut state = self.state.lock().expect("session state not poisoned");
601 let (operations, outgoing, reserved_ids) = match &mut *state {
602 State::Open {
603 operations,
604 outgoing,
605 reserved_ids,
606 ..
607 } => (operations, outgoing, reserved_ids),
608 State::Closed(error) => return Err(error.clone()),
609 };
610 let now = self.now();
611 if now >= operation.deadline {
612 if let OutgoingBody::Reply { id, .. } = body {
613 reserved_ids.remove(&id);
614 }
615 notifications.push(operation.fail(Error::Timeout, now));
616 return Ok(());
617 }
618 let key = OperationKey::new();
619 outgoing.push_back(OutgoingMessage {
620 body,
621 operation: OperationHandle {
622 session: Arc::downgrade(self),
623 key: key.clone(),
624 },
625 #[cfg(any(test, feature = "fuzz"))]
626 deadline: operation.deadline,
627 });
628 operations.insert(key, operation);
629 }
630
631 self.changed.notify_all();
633 Ok(())
634 }
635
636 pub(super) fn expire(&self) {
640 let mut notifications = Notifications::default();
642 let mut state = self.state.lock().expect("session state not poisoned");
643 state.expire(self.now(), &mut notifications);
644 }
645
646 #[cfg(any(test, feature = "fuzz"))]
648 pub(super) fn next_deadline(&self) -> Option<Instant> {
649 let state = self.state.lock().expect("session state not poisoned");
650 match &*state {
651 State::Open { operations, .. } => operations
652 .values()
653 .map(|operation| operation.deadline)
654 .min(),
655 State::Closed(_) => None,
656 }
657 }
658
659 #[cfg(any(test, feature = "fuzz"))]
661 pub(super) fn take_outgoing(&self) -> Option<OutgoingMessage> {
662 let mut notifications = Notifications::default();
664 let mut state = self.state.lock().expect("session state not poisoned");
665 state.expire(self.now(), &mut notifications);
666
667 match &mut *state {
669 State::Open {
670 outgoing,
671 reserved_ids,
672 ..
673 } => {
674 let message = outgoing.pop_front()?;
675 if let OutgoingBody::Reply { id, .. } = &message.body {
676 reserved_ids.remove(id);
677 }
678 Some(message)
679 }
680 State::Closed(_) => None,
681 }
682 }
683
684 pub(super) fn record_write(&self, key: &OperationKey, result: Result<(), Error>) {
687 let mut notifications = Notifications::default();
689 let mut state = self.state.lock().expect("session state not poisoned");
690 let State::Open { operations, .. } = &mut *state else {
691 return;
692 };
693 let Some(operation) = operations.get(key) else {
694 return;
695 };
696
697 let now = self.now();
699 if now >= operation.deadline || result.is_err() {
700 let operation = operations.remove(key).expect("operation held under lock");
701 notifications.push(operation.fail(result.err().unwrap_or(Error::Timeout), now));
702 } else if !operation.sender.response {
703 let operation = operations.remove(key).expect("operation held under lock");
704 notifications.push(operation.sender.send(Ok(PromiseResult::Written)));
705 }
706 }
707
708 #[cfg(any(test, feature = "fuzz"))]
711 pub(super) fn record_response(
712 self: &Arc<Self>,
713 key: &OperationKey,
714 result: Result<Message, Error>,
715 ) {
716 let mut notifications = Notifications::default();
718 let mut state = self.state.lock().expect("session state not poisoned");
719 let State::Open {
720 operations,
721 max_inbound_bytes,
722 ..
723 } = &mut *state
724 else {
725 return;
726 };
727 let Some(operation) = operations.remove(key) else {
728 return;
729 };
730 let result = match result {
731 Ok(message) => Ok(message),
732 Err(Error::Remote(error)) => Err(error),
733 Err(error) => {
734 notifications.push(operation.fail(error, self.now()));
735 return;
736 }
737 };
738
739 let peer = match self.side {
741 Side::Client => Side::Server,
742 Side::Server => Side::Client,
743 };
744 let bytes = peer
745 .encode(0, result)
746 .expect("fixture response belongs to the peer");
747 let bytes = Bytes::from(bytes.into_boxed_slice());
748 let header = self
749 .side
750 .decode_header(bytes.clone())
751 .expect("fixture response has a valid envelope");
752 let result = operation.complete_response(self.now(), &mut notifications, || {
753 self.retain_incoming(bytes, header, *max_inbound_bytes)
754 });
755
756 drop(state);
758 if let Err(error) = result {
759 self.close(error);
760 }
761 }
762
763 fn now(&self) -> Instant {
766 self.clock.now()
767 }
768
769 pub(super) fn handle_message(self: &Arc<Self>, bytes: Vec<u8>) -> Result<(), Error> {
773 let bytes = Bytes::from(bytes.into_boxed_slice());
775 let header = self.side.decode_header(bytes.clone())?;
776 let mut unknown = None;
777
778 let mut notifications = Notifications::default();
780 {
781 let mut state = self.state.lock().expect("session state not poisoned");
782 match MessageKind::from_id(header.id, self.side.into()) {
783 MessageKind::Request => {
784 if header.failed {
785 return Err(self.side.malformed(
786 Some(header),
787 bytes.len(),
788 "envelope",
789 "request contains an error",
790 ));
791 }
792 self.admit_request(&mut state, header, bytes)?;
795 if header.unknown {
796 unknown = Some(header.id);
797 }
798 }
799 MessageKind::Response => {
800 let State::Open {
801 operations,
802 outstanding,
803 max_inbound_bytes,
804 ..
805 } = &mut *state
806 else {
807 let State::Closed(error) = &*state else {
808 unreachable!()
809 };
810 return Err(error.clone());
811 };
812 if let Some(key) = outstanding.remove(&header.id) {
815 if let Some(operation) = operations.remove(&key) {
816 tracing::trace!(
817 "received response {} ({})",
818 header.id,
819 header.payload.unwrap_or("none")
820 );
821 operation.complete_response(self.now(), &mut notifications, || {
822 self.retain_incoming(bytes, header, *max_inbound_bytes)
823 })?;
824 } else {
825 tracing::debug!("discarding late response {}", header.id);
826 }
827 } else {
828 tracing::warn!("discarding unmatched response {}", header.id);
829 }
830 }
831 }
832 }
833
834 if let Some(id) = unknown {
836 self.reply_unknown(id);
837 }
838 self.changed.notify_all();
839 Ok(())
840 }
841
842 fn admit_request(
846 self: &Arc<Self>,
847 state: &mut State,
848 header: Header,
849 bytes: Bytes,
850 ) -> Result<(), Error> {
851 let State::Open {
852 incoming,
853 reserved_ids,
854 max_inbound_requests,
855 max_inbound_bytes,
856 ..
857 } = state
858 else {
859 let State::Closed(error) = state else {
860 unreachable!()
861 };
862 return Err(error.clone());
863 };
864 if reserved_ids.contains(&header.id) {
865 return Err(self.side.malformed(
866 Some(header),
867 bytes.len(),
868 "envelope",
869 "duplicate request ID",
870 ));
871 }
872 let id = header.id;
873 if reserved_ids.len() >= *max_inbound_requests {
874 tracing::warn!(
875 "inbound request limit exceeded (id: {}, used: {}, limit: {})",
876 id,
877 reserved_ids.len(),
878 max_inbound_requests
879 );
880 return Err(Error::InboundRequestLimitExceeded(*max_inbound_requests));
881 }
882 let payload = header.payload.unwrap_or("none");
883 let message = if header.unknown {
884 None
885 } else {
886 Some(self.retain_incoming(bytes, header, *max_inbound_bytes)?)
887 };
888 reserved_ids.insert(id);
889 if let Some(message) = message {
890 incoming.push_back((id, message));
891 }
892 tracing::trace!("received request {} ({})", id, payload);
893 Ok(())
894 }
895
896 pub(super) fn next_outgoing(&self) -> Option<(u64, OutgoingMessage)> {
901 let mut notifications = Notifications::default();
903 let mut state = self.state.lock().expect("session state not poisoned");
904 loop {
905 state.expire(self.now(), &mut notifications);
907 if !notifications.is_empty() {
908 drop(state);
909 drop(notifications);
910 notifications = Notifications::default();
911 state = self.state.lock().expect("session state not poisoned");
912 continue;
913 }
914
915 let State::Open {
917 outgoing,
918 next_id,
919 outstanding,
920 operations,
921 reserved_ids,
922 ..
923 } = &mut *state
924 else {
925 return None;
926 };
927 if let Some(outgoing) = outgoing.pop_front() {
928 let id = match &outgoing.body {
929 OutgoingBody::Request(_) => {
930 let id = next_id.expect("wire request IDs exhausted");
931 *next_id = id.checked_add(2);
932 outstanding.insert(id, outgoing.operation.key.clone());
935 if let Some(operation) = operations.get_mut(&outgoing.operation.key) {
936 operation.log_id = Some(LogId::from(id));
937 }
938 id
939 }
940 OutgoingBody::Reply { id, .. } => {
941 reserved_ids.remove(id);
945 *id
946 }
947 };
948 return Some((id, outgoing));
949 }
950
951 state = self
953 .changed
954 .wait(state)
955 .expect("session state not poisoned");
956 }
957 }
958
959 fn run_writer(&self, sender: transport::Sender<impl Write>) {
962 while let Some((id, outgoing)) = self.next_outgoing() {
963 let request = matches!(outgoing.body, OutgoingBody::Request(_));
966 let kind = if request { "request" } else { "reply" };
967 let body = match outgoing.body {
968 OutgoingBody::Request(body) => Ok(body),
969 OutgoingBody::Reply { result, .. } => result,
970 };
971 let payload = match &body {
972 Ok(message) => message.field_name(),
973 Err(_) => "err",
974 };
975 let result = self.side.encode(id, body).and_then(|bytes| {
976 tracing::trace!("sending {} {} ({})", kind, id, payload);
979 sender.send(&bytes).map_err(Error::from)
980 });
981 match &result {
982 Ok(()) => {}
983 Err(Error::Transport(error)) => {
984 self.close(Error::Transport(error.clone()));
987 break;
988 }
989 Err(error) => tracing::debug!("not sending {} {}: {}", kind, id, error),
990 }
991 {
992 let mut state = self.state.lock().expect("session state not poisoned");
995 if let State::Open { outstanding, .. } = &mut *state
996 && request
997 && result.is_err()
998 {
999 outstanding.remove(&id);
1000 }
1001 }
1002 outgoing.operation.record_write(result);
1005 }
1006 #[cfg(any(test, feature = "fuzz"))]
1007 if let Some((entered, released)) = self.disconnect_hook.lock().unwrap().take() {
1008 let _ = entered.send(());
1009 let _ = released.recv();
1010 }
1011 if let Err(error) = sender.disconnect() {
1014 tracing::debug!(
1015 "failed to signal dropped session {}: {}",
1016 sender.log_id(),
1017 error
1018 );
1019 }
1020 }
1021
1022 pub(super) fn run_deadlines(&self) {
1025 let mut notifications = Notifications::default();
1027 let mut state = self.state.lock().expect("session state not poisoned");
1028 loop {
1029 state.expire(self.now(), &mut notifications);
1031 if !notifications.is_empty() {
1032 drop(state);
1033 drop(notifications);
1034 notifications = Notifications::default();
1035 state = self.state.lock().expect("session state not poisoned");
1036 continue;
1037 }
1038 let State::Open { operations, .. } = &*state else {
1039 return;
1040 };
1041
1042 state = match operations
1045 .values()
1046 .map(|operation| operation.deadline)
1047 .min()
1048 {
1049 Some(deadline) => {
1050 self.changed
1051 .wait_deadline(state, deadline)
1052 .expect("session state not poisoned")
1053 .0
1054 }
1055 None => self
1056 .changed
1057 .wait(state)
1058 .expect("session state not poisoned"),
1059 };
1060 }
1061 }
1062}
1063
1064impl Session {
1065 pub(super) fn start<W: Write + Send + 'static>(
1068 side: Side,
1069 clock: Clock,
1070 sender: transport::Sender<W>,
1071 stream_closer: Option<transport::Closer>,
1072 #[cfg(any(test, feature = "fuzz"))] workers: Arc<worker::Tracker>,
1073 ) -> Self {
1074 let session = Self {
1076 inner: Arc::new(SessionInner::new(
1077 side,
1078 clock,
1079 sender.log_id(),
1080 stream_closer,
1081 #[cfg(any(test, feature = "fuzz"))]
1082 workers.clone(),
1083 )),
1084 };
1085
1086 let inner = session.inner.clone();
1088 worker::spawn(
1089 "wire-writer",
1090 #[cfg(any(test, feature = "fuzz"))]
1091 &workers,
1092 move || inner.run_writer(sender),
1093 );
1094
1095 let inner = session.inner.clone();
1097 worker::spawn(
1098 "wire-deadlines",
1099 #[cfg(any(test, feature = "fuzz"))]
1100 &workers,
1101 move || inner.run_deadlines(),
1102 );
1103 session
1104 }
1105}
1106
1107impl State {
1108 fn workload(&self) -> (usize, usize) {
1110 match self {
1111 Self::Open {
1112 operations,
1113 outgoing,
1114 incoming,
1115 ..
1116 } => (operations.len(), outgoing.len() + incoming.len()),
1117 Self::Closed(_) => (0, 0),
1118 }
1119 }
1120
1121 fn close(
1124 &mut self,
1125 error: Error,
1126 now: Instant,
1127 notifications: &mut Notifications,
1128 ) -> Option<Self> {
1129 if let Self::Closed(_) = self {
1131 return None;
1132 }
1133
1134 let mut removed = std::mem::replace(self, Self::Closed(error.clone()));
1136 if let Self::Open { operations, .. } = &mut removed {
1137 for (_, operation) in operations.drain() {
1138 notifications.push(operation.fail(error.clone(), now));
1139 }
1140 }
1141 Some(removed)
1142 }
1143
1144 fn expire(&mut self, now: Instant, notifications: &mut Notifications) {
1147 if let Self::Open {
1148 operations,
1149 outgoing,
1150 reserved_ids,
1151 ..
1152 } = self
1153 {
1154 let expired: Vec<_> = operations
1156 .iter()
1157 .filter(|(_, operation)| now >= operation.deadline)
1158 .map(|(key, _)| key.clone())
1159 .collect();
1160 for key in expired {
1161 notifications.push(
1162 operations
1163 .remove(&key)
1164 .expect("expired operation held under lock")
1165 .fail(Error::Timeout, now),
1166 );
1167 }
1168
1169 outgoing.retain(|outgoing| {
1172 let retained = operations.contains_key(&outgoing.operation.key);
1173 if !retained && let OutgoingBody::Reply { id, .. } = outgoing.body {
1174 reserved_ids.remove(&id);
1175 }
1176 retained
1177 });
1178 }
1179 }
1180}
1181
1182#[cfg(any(test, feature = "fuzz"))]
1184impl Session {
1185 #[cfg(test)]
1187 pub(super) fn fixture() -> Self {
1188 Self::fixture_for(Side::Server)
1189 }
1190
1191 #[cfg(test)]
1193 pub(super) fn fixture_for(side: Side) -> Self {
1194 Self::fixture_with_clock(side, crate::transport::testing::test_clock().clock())
1195 }
1196
1197 pub(super) fn fixture_with_clock(side: Side, clock: Clock) -> Self {
1199 Self {
1200 inner: Arc::new(SessionInner::new(
1201 side,
1202 clock,
1203 LogId::default(),
1204 None,
1205 Arc::new(worker::Tracker::default()),
1206 )),
1207 }
1208 }
1209}
1210
1211#[cfg(any(test, feature = "fuzz"))]
1212impl SessionInner {
1213 pub(super) fn pause_disconnect(
1216 &self,
1217 ) -> (std::sync::mpsc::Receiver<()>, std::sync::mpsc::Sender<()>) {
1218 let (entered, observed) = std::sync::mpsc::channel();
1219 let (release, released) = std::sync::mpsc::channel();
1220 *self.disconnect_hook.lock().unwrap() = Some((entered, released));
1221 (observed, release)
1222 }
1223
1224 pub(super) fn watch_drop(&self) -> std::sync::mpsc::Receiver<()> {
1226 let (sender, receiver) = std::sync::mpsc::channel();
1227 *self.drop_hook.lock().unwrap() = Some(sender);
1228 receiver
1229 }
1230
1231 pub(super) fn use_last_request_id(&self) {
1233 let mut state = self.state.lock().unwrap();
1234 let State::Open {
1235 next_id,
1236 outstanding,
1237 ..
1238 } = &mut *state
1239 else {
1240 panic!("open session required")
1241 };
1242 assert!(outstanding.is_empty());
1243 *next_id = Some(if self.side == Side::Client {
1244 u64::MAX
1245 } else {
1246 u64::MAX - 1
1247 });
1248 }
1249
1250 pub(super) fn wait_response(&self, id: u64) {
1253 let mut state = self.state.lock().unwrap();
1255 while matches!(&*state, State::Open { outstanding, .. } if outstanding.contains_key(&id)) {
1256 state = self.changed.wait(state).unwrap();
1257 }
1258
1259 let State::Open { outstanding, .. } = &*state else {
1261 panic!("session closed before response fence");
1262 };
1263 assert!(
1264 !outstanding.contains_key(&id),
1265 "reader did not process response"
1266 );
1267 }
1268
1269 pub(super) fn outstanding_ids(&self) -> Vec<u64> {
1271 let state = self.state.lock().unwrap();
1272 let State::Open { outstanding, .. } = &*state else {
1273 panic!("open session required")
1274 };
1275 let mut ids: Vec<_> = outstanding.keys().copied().collect();
1276 ids.sort_unstable();
1277 ids
1278 }
1279
1280 pub(super) fn inject_request(self: &Arc<Self>, id: u64, message: Message) -> Result<(), Error> {
1283 let peer = match self.side {
1284 Side::Client => Side::Server,
1285 Side::Server => Side::Client,
1286 };
1287 let bytes = peer
1288 .encode(id, Ok(message))
1289 .expect("fixture request belongs to peer");
1290 let bytes = Bytes::from(bytes.into_boxed_slice());
1291 let header = self.side.decode_header(bytes.clone())?;
1292 let result = {
1293 let mut state = self.state.lock().expect("session state not poisoned");
1294 self.admit_request(&mut state, header, bytes)
1295 };
1296 self.changed.notify_all();
1297 if let Err(error) = &result {
1298 self.close(error.clone());
1299 }
1300 result
1301 }
1302
1303 pub(super) fn inbound_usage(&self) -> (usize, usize) {
1305 let state = self.state.lock().expect("session state not poisoned");
1306 let requests = match &*state {
1307 State::Open { reserved_ids, .. } => reserved_ids.len(),
1308 State::Closed(_) => 0,
1309 };
1310 (requests, self.retained_bytes.load(Ordering::Relaxed))
1311 }
1312
1313 pub(super) fn watch_recv_wait(&self) -> std::sync::mpsc::Receiver<()> {
1320 let (sender, receiver) = std::sync::mpsc::channel();
1321 let mut state = self.state.lock().expect("session state not poisoned");
1322 let State::Open {
1323 incoming,
1324 wait_hook,
1325 ..
1326 } = &mut *state
1327 else {
1328 panic!("only watch an open session receive");
1329 };
1330 assert!(incoming.is_empty());
1331 *wait_hook = Some(sender);
1332 receiver
1333 }
1334}
1335
1336#[cfg(any(test, feature = "fuzz"))]
1337impl Drop for SessionInner {
1338 fn drop(&mut self) {
1340 if let Some(sender) = self.drop_hook.get_mut().unwrap().take() {
1341 let _ = sender.send(());
1342 }
1343 }
1344}
1345
1346#[cfg(test)]
1348#[cfg_attr(coverage_nightly, coverage(off))]
1349mod tests {
1350 use super::{Notifications, SessionInner, Side};
1351 use crate::protocol::{self, Error, Session};
1352 use crate::transport::{Read, Stream, Verifier, Write};
1353 use std::fmt::Debug;
1354
1355 #[test]
1357 fn test_deadline_callback_releases_locks_before_worker_parks() {
1358 use std::sync::mpsc;
1359 use std::thread;
1360 use std::time::Duration;
1361
1362 let mut tester = crate::transport::testing::test_clock();
1364 let session = Session::fixture_with_clock(Side::Client, tester.clock());
1365 let deadline = tester.clock().now() + Duration::from_secs(5);
1366 let mut promise = session.requester().request(vec![1], deadline).unwrap();
1367 let notification = promise.notification_unlocked();
1368 let state = session.inner.clone();
1369 let (observed, receiver) = mpsc::channel();
1370 promise.notify(move || {
1371 let _ = observed.send((state.state.try_lock().is_ok(), notification()));
1372 });
1373 let state = session.inner.clone();
1374 let worker = thread::spawn(move || state.run_deadlines());
1375
1376 tester.wait_blocked(1);
1378 tester.advance_to(deadline);
1379 let unlocked = receiver.recv().unwrap();
1380
1381 session.close();
1383 worker.join().unwrap();
1384 assert_eq!(unlocked, (true, true));
1385 assert!(matches!(promise.wait::<Vec<u8>>(), Err(Error::Timeout)));
1386 }
1387
1388 #[test]
1390 fn test_notification_releases_locks_on_every_settlement_path() {
1391 use crate::protocol::Promise;
1392 use std::sync::{Arc, mpsc};
1393 use std::thread::{self, ThreadId};
1394 use std::time::Duration;
1395
1396 struct Capture {
1398 session: Arc<SessionInner>,
1400 notification: Box<dyn Fn() -> bool + Send>,
1402 observed: mpsc::Sender<(bool, bool, ThreadId)>,
1404 }
1405
1406 impl Capture {
1407 fn check(&self) {
1409 let _ = self.observed.send((
1410 self.session.state.try_lock().is_ok(),
1411 (self.notification)(),
1412 thread::current().id(),
1413 ));
1414 }
1415 }
1416
1417 impl Drop for Capture {
1418 fn drop(&mut self) {
1420 self.check();
1421 }
1422 }
1423
1424 fn check<T>(
1426 mut promise: Promise<T>,
1427 session: Arc<SessionInner>,
1428 completed: bool,
1429 path: &str,
1430 settle: impl FnOnce() + Send,
1431 ) {
1432 let (observed, receiver) = mpsc::channel();
1434 let capture = Capture {
1435 session,
1436 notification: Box::new(promise.notification_unlocked()),
1437 observed,
1438 };
1439
1440 let expected = thread::scope(|scope| {
1442 if completed {
1443 scope.spawn(settle).join().unwrap();
1444 promise.notify(move || capture.check());
1445 thread::current().id()
1446 } else {
1447 promise.notify(move || capture.check());
1448 let settling = scope.spawn(move || {
1449 settle();
1450 thread::current().id()
1451 });
1452 settling.join().unwrap()
1453 }
1454 });
1455
1456 let observations: Vec<_> = receiver.try_iter().collect();
1458 assert_eq!(
1459 observations,
1460 vec![(true, true, expected); 2],
1461 "{path}, completed={completed}"
1462 );
1463 }
1464
1465 for path in [
1467 "response",
1468 "response limit",
1469 "late response",
1470 "fixture response",
1471 "fixture error",
1472 "write failure",
1473 "late write",
1474 "reply",
1475 "close",
1476 "expire",
1477 "take outgoing",
1478 "next outgoing",
1479 "limits",
1480 ] {
1481 for completed in [false, true] {
1482 let mut tester = crate::transport::testing::test_clock();
1484 let clock = tester.clock();
1485 let deadline = clock.now() + Duration::from_secs(5);
1486 let mut session = Session::fixture_with_clock(Side::Client, clock);
1487 let inner = session.inner.clone();
1488
1489 if path == "reply" {
1491 inner.inject_request(2, vec![1].into()).unwrap();
1492 let (_, responder) = session.recv().unwrap();
1493 let promise = responder.reply(vec![2], deadline).unwrap();
1494 let outgoing = inner.take_outgoing().unwrap();
1495 check(promise, inner, completed, path, move || {
1496 outgoing.operation.record_write(Ok(()))
1497 });
1498 continue;
1499 }
1500
1501 if path == "response limit" {
1503 session = session.set_inbound_limits(1, 0);
1504 }
1505 if path == "limits" {
1506 inner.inject_request(2, vec![1].into()).unwrap();
1507 }
1508 let promise = session.requester().request(vec![1], deadline).unwrap();
1509 let (id, outgoing) = inner.next_outgoing().unwrap();
1510 let next = if path == "next outgoing" {
1511 Some(
1512 session
1513 .requester()
1514 .request(vec![2], deadline + Duration::from_secs(5))
1515 .unwrap(),
1516 )
1517 } else {
1518 None
1519 };
1520
1521 check(promise, inner.clone(), completed, path, || match path {
1523 "response" | "response limit" | "late response" => {
1524 if path == "late response" {
1526 tester.advance_to(deadline);
1527 }
1528
1529 let bytes = Side::Server.encode(id, Ok(vec![2].into())).unwrap();
1531 let result = inner.handle_message(bytes);
1532 assert_eq!(result.is_ok(), path != "response limit");
1533 }
1534 "fixture response" => outgoing.operation.record_response(Ok(vec![2].into())),
1535 "fixture error" => {
1536 inner.record_response(&outgoing.operation.key, Err(Error::Closed))
1537 }
1538 "write failure" => outgoing.operation.record_write(Err(Error::Closed)),
1539 "late write" => {
1540 tester.advance_to(deadline);
1541 outgoing.operation.record_write(Ok(()));
1542 }
1543 "close" => inner.close(Error::Closed),
1544 "limits" => {
1545 let mut notifications = Notifications::default();
1546 inner.set_inbound_limits(0, 1024, &mut notifications);
1547 }
1548 "expire" | "take outgoing" | "next outgoing" => {
1549 tester.advance_to(deadline);
1550 match path {
1551 "expire" => inner.expire(),
1552 "take outgoing" => assert!(inner.take_outgoing().is_none()),
1553 "next outgoing" => assert!(inner.next_outgoing().is_some()),
1554 _ => unreachable!(),
1555 }
1556 }
1557 _ => unreachable!(),
1558 });
1559 drop(next);
1560 }
1561 }
1562 }
1563
1564 #[allow(dead_code)]
1566 fn connect<R, W, V>(stream: Stream<R, W>, verifier: &V) -> Result<(Session, V::Info), Error>
1567 where
1568 R: Read + Send + 'static,
1569 W: Write + Send + 'static,
1570 V: Verifier,
1571 {
1572 protocol::connect(stream, verifier)
1573 }
1574
1575 #[test]
1578 fn test_thread_capabilities() {
1579 fn movable<T: Debug + Send + 'static>() {}
1581 movable::<Session>();
1582 }
1583}