ruststream-zeromq 0.7.1

ZeroMQ transport implementation of the RustStream broker contract, for bridging to non-Rust peers.
Documentation
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
906
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
932
933
934
935
936
937
938
939
940
941
942
943
944
945
946
947
948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993
994
995
996
//! Machinery shared by the three socket patterns.

use std::fmt;
use std::future::{Future, poll_fn};
use std::mem;
use std::num::NonZeroUsize;
use std::panic;
use std::pin::Pin;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex as StdMutex, PoisonError, Weak};
use std::task::{Context, Poll, ready};
use std::time::Duration;

use aliasable::boxed::AliasableBox;
use futures::Stream;
use futures::lock::Mutex;
use ruststream::{Subscriber, nonzero};
use selfie::SelfieMut;
use selfie::refs::RefType;
use tokio::runtime::Handle;
use tokio::sync::{OnceCell, mpsc};
use tokio::task::JoinHandle;
use zeromq::prelude::*;
use zeromq::{Socket, ZmqError as WireError, ZmqResult};

use crate::endpoint::{Endpoint, Side};
use crate::error::{ZmqError, box_err};
use crate::fanout::PubSide;
#[cfg(feature = "testing")]
use crate::in_process::{Bus, Deliveries, Registration};
use crate::message::ZmqMessage;
use crate::queue::Attached as QueueAttached;

/// How long a send retries while the ZMTP handshake settles: the implementation returns the
/// message immediately when no peer is attached yet, which is routine right after `connect`.
pub(crate) const SEND_RETRY_WINDOW: Duration = Duration::from_secs(5);
pub(crate) const SEND_RETRY_STEP: Duration = Duration::from_millis(50);

/// How long a partial batch waits for more deliveries after its first one.
///
/// ZMTP carries one multipart message per receive, so batches are assembled on the client and this
/// deadline is the crate's own choice - the batch size is not, it arrives per subscription. Twenty
/// milliseconds coalesces a burst that is already queued behind the socket while costing an idle
/// subscription far less than the round trip it is waiting on anyway.
pub(crate) const BATCH_MAX_WAIT: Duration = Duration::from_millis(20);

/// How many deliveries a subscription reads off its socket ahead of the handler, unless its
/// broker names another bound: the receive high-water mark `ZeroMQ` itself gives a socket.
pub(crate) const DEFAULT_READ_AHEAD: NonZeroUsize = nonzero!(1000_usize);

/// The channel between a subscription's driver task and its handler, `read_ahead` deep.
///
/// Bounded, so a full channel stops the driver, the driver stops reading the socket, and the
/// socket holds the sender back: a handler slower than the wire slows the wire down instead of
/// accumulating the difference in memory. The bound is the connected form's own, taken from the
/// broker it was connected from, so clones sharing one lifecycle keep their own bounds.
pub(crate) fn delivery_channel(read_ahead: NonZeroUsize) -> (DeliverySender, DeliveryReceiver) {
    let (tx, rx) = mpsc::channel(read_ahead.get());
    (tx, DeliveryReceiver::Driver(rx))
}

/// What a broker connected in the test harness runs over: the ZMTP sockets `connect` attaches,
/// or the in-process transport `connect_in_process` gives it instead.
///
/// It exists only under the `testing` feature: a production lifecycle has no field for it, so it
/// carries no second transport and no branch to one.
#[cfg(feature = "testing")]
#[derive(Debug)]
pub(crate) enum Transport {
    Sockets,
    InProcess(Arc<Bus>),
}

/// Shared lifecycle state: the endpoint, the listener a local subscription bound (which is what a
/// same-process publisher dials for the loopback arrangement), the subscription that dialed the
/// endpoint on the other side, and the closed flag aliased handles trip over.
#[derive(Debug)]
pub(crate) struct Lifecycle {
    pub(crate) endpoint: Endpoint,
    #[cfg(feature = "testing")]
    transport: Transport,
    /// The listener of the subscription bound here now; empty again once it has gone, so a later
    /// subscription can bind the endpoint.
    listener: Arc<StdMutex<Option<Arc<Listener>>>>,
    /// Runs one bind at a time, so two subscriptions opening together cannot both bind, and
    /// orders each publisher attach against them, so a subscription that binds finds every
    /// publisher listening here in [`listening`](Self::listening).
    binding: tokio::sync::Mutex<()>,
    /// The publishers of this service that listen on the endpoint themselves, because no
    /// subscription held it when they attached. A subscription that binds the endpoint later
    /// makes them give it up first, so their next publish dials that subscription.
    listening: StdMutex<Vec<OwnListener>>,
    /// The first subscription of this service that dialed the endpoint. The peer it reads from is
    /// the sending end of the pattern, which takes nothing, so a publisher of the same broker has
    /// nowhere to send.
    dialer: OnceCell<String>,
    pub(crate) closed: AtomicBool,
    /// The runtime the broker connected on. A socket registers its connection with the runtime
    /// that binds or dials it and starts its accept loop and peer tasks there, so every socket
    /// attaches here and every driver task runs here, whichever runtime asked: a handler on a
    /// thread of its own publishes from a runtime that stops before the broker does.
    runtime: Handle,
}

/// The socket a subscription in this process bound on the endpoint.
#[derive(Debug)]
pub(crate) struct Listener {
    /// The address the bind resolved to: the port the operating system chose for a `:0`
    /// endpoint.
    pub(crate) address: String,
    /// The subscription that bound it, which is what receives a message sent there.
    pub(crate) subscription: String,
}

/// Where a sending socket attaches.
#[derive(Debug, Clone)]
pub(crate) enum SendTarget<'a> {
    /// Another process listens on the endpoint, so the socket dials it.
    Dial(&'a str),
    /// Nothing in this process holds the endpoint, so the socket listens there and peers dial it.
    Listen(&'a str),
    /// A subscription in this process bound the endpoint, so the socket dials the address that
    /// bind resolved to.
    Local(Arc<Listener>),
}

impl SendTarget<'_> {
    /// The address the socket attaches to.
    pub(crate) fn address(&self) -> &str {
        match self {
            Self::Dial(address) | Self::Listen(address) => address,
            Self::Local(listener) => &listener.address,
        }
    }
}

impl Lifecycle {
    pub(crate) fn new(endpoint: Endpoint) -> Self {
        Self {
            endpoint,
            #[cfg(feature = "testing")]
            transport: Transport::Sockets,
            listener: Arc::new(StdMutex::new(None)),
            binding: tokio::sync::Mutex::new(()),
            listening: StdMutex::new(Vec::new()),
            dialer: OnceCell::new(),
            closed: AtomicBool::new(false),
            // Built inside `connect`, which the runtime the broker connects on polls.
            runtime: Handle::current(),
        }
    }

    /// The lifecycle of a broker the test harness connected in process: the same endpoint and
    /// the same bookkeeping, over a transport of its own.
    #[cfg(feature = "testing")]
    pub(crate) fn in_process(endpoint: Endpoint) -> Self {
        Self {
            transport: Transport::InProcess(Arc::new(Bus::default())),
            ..Self::new(endpoint)
        }
    }

    /// The in-process transport, when the test harness connected this broker in process.
    #[cfg(feature = "testing")]
    pub(crate) fn in_process_bus(&self) -> Option<&Arc<Bus>> {
        match &self.transport {
            Transport::Sockets => None,
            Transport::InProcess(bus) => Some(bus),
        }
    }

    /// Whether this broker dials the endpoint, so a sending socket's only peer is the one at
    /// that address.
    pub(crate) fn dials(&self) -> bool {
        self.endpoint.side() == Side::Connect
    }

    /// The subscription of this service that dialed the endpoint, or `None` until one has.
    pub(crate) fn dialer(&self) -> Option<&str> {
        self.dialer.get().map(String::as_str)
    }

    /// The address a local subscription resolved by binding, or `None` until one has bound.
    pub(crate) fn bound_address(&self) -> Option<String> {
        self.listener().map(|listener| listener.address.clone())
    }

    fn listener(&self) -> Option<Arc<Listener>> {
        self.listener
            .lock()
            .unwrap_or_else(PoisonError::into_inner)
            .clone()
    }

    /// Runs `attach` on the runtime the broker connected on and returns its outcome: one spawn
    /// per subscription or publisher attach, never per message.
    pub(crate) async fn on_runtime<Attach, Output>(
        &self,
        attach: Attach,
    ) -> Result<Output, ZmqError>
    where
        Attach: Future<Output = Result<Output, ZmqError>> + Send + 'static,
        Output: Send + 'static,
    {
        match self.runtime.spawn(attach).await {
            Ok(outcome) => outcome,
            Err(err) if err.is_panic() => panic::resume_unwind(err.into_panic()),
            // Cancelled only by that runtime shutting down, which takes the broker with it.
            Err(_) => Err(ZmqError::NotConnected),
        }
    }

    /// Starts a subscription's driver task on the runtime the broker connected on.
    pub(crate) fn spawn<Task>(&self, task: Task) -> JoinHandle<()>
    where
        Task: Future<Output = ()> + Send + 'static,
    {
        self.runtime.spawn(task)
    }

    pub(crate) fn ensure_open(&self) -> Result<(), ZmqError> {
        if self.closed.load(Ordering::Acquire) {
            return Err(ZmqError::NotConnected);
        }
        Ok(())
    }

    /// Attaches a receiving socket for `subscription` per the endpoint's side, recording the
    /// listener on bind so a same-process publisher can dial it, and the dialing subscription on
    /// connect so a same-process publisher is refused before it dials a peer that only sends.
    ///
    /// On the bind side the answer is the slot the subscription holds: the subscription keeps it
    /// for as long as it is open, and dropping it frees the endpoint for the next one.
    pub(crate) async fn attach_receiver<S: Socket>(
        &self,
        socket: &mut S,
        subscription: &str,
    ) -> Result<Option<BoundSlot>, ZmqError> {
        match self.endpoint.side() {
            Side::Bind => self
                .hold_listener(subscription, async || {
                    socket
                        .bind(self.endpoint.address())
                        .await
                        .map(|resolved| resolved.to_string())
                        .map_err(|e| ZmqError::Endpoint {
                            endpoint: self.endpoint.address().to_owned(),
                            source: box_err(e),
                        })
                })
                .await
                .map(Some),
            Side::Connect => {
                socket
                    .connect(self.endpoint.address())
                    .await
                    .map_err(|e| ZmqError::Endpoint {
                        endpoint: self.endpoint.address().to_owned(),
                        source: box_err(e),
                    })?;
                let _ = self.dialer.set(subscription.to_owned());
                Ok(None)
            }
        }
    }

    /// Takes the endpoint for `subscription` on the in-process transport, with the bookkeeping
    /// [`attach_receiver`](Self::attach_receiver) keeps: the first subscription holds a bound
    /// endpoint and a second is refused, and a subscription that dials is recorded so a publisher
    /// of this service is refused in the socket's words. There is no socket, so the listener's
    /// address is the endpoint as configured, and the bound subscription holds its slot as a
    /// socket does, until it closes.
    #[cfg(feature = "testing")]
    pub(crate) async fn attach_in_process(
        &self,
        subscription: &str,
    ) -> Result<Option<BoundSlot>, ZmqError> {
        match self.endpoint.side() {
            Side::Bind => self
                .hold_listener(subscription, async || {
                    Ok(self.endpoint.address().to_owned())
                })
                .await
                .map(Some),
            Side::Connect => {
                let _ = self.dialer.set(subscription.to_owned());
                Ok(None)
            }
        }
    }

    /// Binds the endpoint for `subscription` through `bind`: the subscription holds the endpoint
    /// until the returned slot is dropped, and one that opens while another holds it is refused.
    async fn hold_listener<Bind, Bound>(
        &self,
        subscription: &str,
        bind: Bind,
    ) -> Result<BoundSlot, ZmqError>
    where
        Bind: FnOnce() -> Bound,
        Bound: Future<Output = Result<String, ZmqError>>,
    {
        // One bind at a time: a subscription opening while another holds the endpoint finds its
        // listener and is refused. A bind that fails leaves the slot empty.
        let _binding = self.binding.lock().await;
        // Refused at startup rather than by the type: how many subscriptions a scope mounts on
        // one broker is decided by statements the type does not count.
        if let Some(held) = self.listener() {
            return Err(endpoint_taken(subscription, &held.subscription));
        }
        // A publisher that attached while the endpoint was free listens on it itself: on a fixed
        // port the bind below would find the address taken, and on port zero the publisher would
        // keep sending to its own listener, which nothing reads.
        let listening = mem::take(
            &mut *self
                .listening
                .lock()
                .unwrap_or_else(PoisonError::into_inner),
        );
        for publisher in listening {
            publisher.give_way().await;
        }
        let listener = Arc::new(Listener {
            address: bind().await?,
            subscription: subscription.to_owned(),
        });
        *self.listener.lock().unwrap_or_else(PoisonError::into_inner) = Some(Arc::clone(&listener));
        Ok(BoundSlot {
            slot: Arc::clone(&self.listener),
            listener,
        })
    }

    /// Where a sending socket attaches: the endpoint itself when dialing out or when nothing in
    /// this process holds it, or the listener a subscription in this process bound (the loopback
    /// arrangement).
    pub(crate) fn send_target(&self) -> SendTarget<'_> {
        match self.endpoint.side() {
            Side::Connect => SendTarget::Dial(self.endpoint.address()),
            Side::Bind => self.listener().map_or_else(
                || SendTarget::Listen(self.endpoint.address()),
                SendTarget::Local,
            ),
        }
    }

    /// The listener a sending socket of this service would dial, on the in-process transport:
    /// what [`attach_sender`](Self::attach_sender) reports, with no socket to attach.
    #[cfg(feature = "testing")]
    pub(crate) fn local_listener(&self) -> Option<Arc<Listener>> {
        match self.send_target() {
            SendTarget::Local(listener) => Some(listener),
            SendTarget::Dial(_) | SendTarget::Listen(_) => None,
        }
    }

    /// Holds off subscriptions binding the endpoint while a publisher attaches, so the place the
    /// publisher picks stays the place it attaches to, and a publisher that listens is on
    /// [`listens_alongside`](Self::listens_alongside)'s list before a subscription looks.
    ///
    /// Taken before the publisher's own slot, as a subscription takes them, and only on the
    /// attach, never on a send.
    pub(crate) async fn attaching(&self) -> tokio::sync::MutexGuard<'_, ()> {
        self.binding.lock().await
    }

    /// Whether a sending socket attached on the bind side with no subscription of this service
    /// holding the endpoint, and therefore listens on the endpoint itself.
    pub(crate) fn listens_here(&self, dialed_local: bool) -> bool {
        self.endpoint.side() == Side::Bind && !dialed_local
    }

    /// Whether a publisher attaches its socket when it is paired rather than on its first
    /// publish: on the bind side, so the socket listens from the moment the service starts and a
    /// peer that dials the endpoint finds it before anything is sent. A dialing publisher waits
    /// for its first publish, because the peer it dials may not be up yet.
    pub(crate) fn attaches_on_pairing(&self) -> bool {
        self.endpoint.side() == Side::Bind && !self.closed.load(Ordering::Acquire)
    }

    /// Records a publisher listening on the endpoint itself, for a subscription that binds the
    /// endpoint later to make it give way. Called under [`attaching`](Self::attaching).
    pub(crate) fn listens_alongside(&self, publisher: OwnListener) {
        let mut listening = self
            .listening
            .lock()
            .unwrap_or_else(PoisonError::into_inner);
        listening.retain(OwnListener::is_alive);
        listening.push(publisher);
    }

    /// Attaches a sending socket per [`send_target`](Self::send_target), and returns the local
    /// listener it dialed, if it dialed one.
    pub(crate) async fn attach_sender<S: Socket>(
        &self,
        socket: &mut S,
    ) -> Result<Option<Arc<Listener>>, ZmqError> {
        attach_to(socket, self.send_target()).await
    }
}

/// Attaches a sending socket to `target`, a place [`Lifecycle::send_target`] named, and returns
/// the local listener it dialed, if it dialed one.
pub(crate) async fn attach_to<S: Socket>(
    socket: &mut S,
    target: SendTarget<'_>,
) -> Result<Option<Arc<Listener>>, ZmqError> {
    let outcome = match &target {
        SendTarget::Listen(address) => socket.bind(address).await.map(|_| ()),
        SendTarget::Dial(address) => socket.connect(address).await,
        SendTarget::Local(listener) => socket.connect(&listener.address).await,
    };
    outcome.map_err(|e| ZmqError::Endpoint {
        endpoint: target.address().to_owned(),
        source: box_err(e),
    })?;
    Ok(match target {
        SendTarget::Local(listener) => Some(listener),
        SendTarget::Dial(_) | SendTarget::Listen(_) => None,
    })
}

/// A publisher slot of this service whose socket listens on the endpoint itself.
pub(crate) enum OwnListener {
    Queue(Weak<Mutex<Option<QueueAttached>>>),
    Fanout(Weak<Mutex<Option<Sender<PubSide>>>>),
}

impl fmt::Debug for OwnListener {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.write_str(match self {
            Self::Queue(_) => "OwnListener::Queue",
            Self::Fanout(_) => "OwnListener::Fanout",
        })
    }
}

impl OwnListener {
    fn is_alive(&self) -> bool {
        match self {
            Self::Queue(slot) => slot.strong_count() > 0,
            Self::Fanout(slot) => slot.strong_count() > 0,
        }
    }

    /// Closes the publisher's listener and empties its slot, so its next publish attaches again,
    /// to the subscription about to bind. Waits for a send in flight on it to finish first.
    async fn give_way(self) {
        match self {
            Self::Queue(slot) => {
                if let Some(slot) = slot.upgrade() {
                    let attached = slot.lock().await.take();
                    if let Some(attached) = attached {
                        attached.give_way().await;
                    }
                }
            }
            Self::Fanout(slot) => {
                if let Some(slot) = slot.upgrade() {
                    let sender = slot.lock().await.take();
                    if let Some(Sender::Socket(PubSide::Remote(mut socket))) = sender {
                        release(&mut socket).await;
                    }
                }
            }
        }
    }
}

/// Closes every listener `socket` holds and waits until each has let its address go, so a
/// subscription can bind it next.
pub(crate) async fn release<S: Socket>(socket: &mut S) {
    for err in socket.unbind_all().await {
        tracing::warn!(error = %err, "a publisher's listener did not close cleanly");
    }
}

/// The endpoint a bound subscription holds: dropping it, when the subscription is dropped,
/// frees the endpoint for the next subscription, unless another has taken it since.
#[derive(Debug)]
pub(crate) struct BoundSlot {
    slot: Arc<StdMutex<Option<Arc<Listener>>>>,
    listener: Arc<Listener>,
}

impl Drop for BoundSlot {
    fn drop(&mut self) {
        let mut slot = self.slot.lock().unwrap_or_else(PoisonError::into_inner);
        if slot
            .as_ref()
            .is_some_and(|held| Arc::ptr_eq(held, &self.listener))
        {
            *slot = None;
        }
    }
}

/// The refusal of a queue publish that would come back to this service: the subscription that
/// bound the queue receives every message pushed into it, whatever its name, so a message named
/// anything else would arrive there as its next delivery.
pub(crate) fn returns_to_subscription(name: &str, subscription: &str) -> ZmqError {
    ZmqError::Send {
        name: name.to_owned(),
        reason: format!(
            "this service's subscription '{subscription}' binds the queue and receives every \
             message pushed into it; publish '{name}' through a queue on an endpoint of its own"
        ),
    }
}

/// The refusal of a second subscription on an endpoint this service binds: the endpoint is one
/// listening socket, read by the subscription that bound it, and another subscription would bind
/// a socket of its own that nothing dials (on an ephemeral port) or fail to bind (on a fixed one).
pub(crate) fn endpoint_taken(subscription: &str, holder: &str) -> ZmqError {
    ZmqError::Invalid(format!(
        "subscription '{subscription}' cannot open on the endpoint this broker binds: this \
         service's subscription '{holder}' binds it, and a bound endpoint is one socket read by \
         the subscription that bound it; mount '{subscription}' on a broker with an endpoint of \
         its own"
    ))
}

/// The refusal of a publish on an endpoint a subscription of this service dials: the peer that
/// subscription reads from is the sending end of the pattern (`sender`, a PUSH or a PUB socket),
/// and the handshake refuses a sending socket that dials it.
///
/// The publishers refuse in these words before they dial.
pub(crate) fn dials_the_sender(
    name: &str,
    subscription: &str,
    sender: &str,
    pattern: &str,
) -> ZmqError {
    ZmqError::Send {
        name: name.to_owned(),
        reason: format!(
            "this service's subscription '{subscription}' dials the endpoint, and the {sender} \
             socket it reads from there takes no message; publish '{name}' through a {pattern} \
             on an endpoint of its own"
        ),
    }
}

/// Sends with a bounded retry while the handshake settles; `ReturnToSender` hands the message
/// back, so nothing is lost by retrying.
pub(crate) async fn send_with_retry<S: SocketSend>(
    socket: &mut S,
    name: &str,
    message: zeromq::ZmqMessage,
) -> Result<(), ZmqError> {
    retry(Plain(socket), message)
        .await
        .map_err(|failure| failure.into_error(name))
}

/// One attempt at handing a message to a socket, which [`retry`] repeats while the handshake
/// settles.
trait Offer {
    fn offer(&mut self, message: zeromq::ZmqMessage) -> impl Future<Output = ZmqResult<()>> + Send;
}

/// A socket offered a message directly, in the caller's future.
struct Plain<'a, S>(&'a mut S);

impl<S: SocketSend> Offer for Plain<'_, S> {
    fn offer(&mut self, message: zeromq::ZmqMessage) -> impl Future<Output = ZmqResult<()>> + Send {
        self.0.send(message)
    }
}

/// Why a send did not leave, before it is reported against the name it carried.
#[derive(Debug)]
pub(crate) enum SendFailure {
    /// No peer attached within the retry window.
    NoPeer,
    /// The socket refused the message for another reason.
    Wire(WireError),
}

impl SendFailure {
    pub(crate) fn into_error(self, name: &str) -> ZmqError {
        ZmqError::Send {
            name: name.to_owned(),
            reason: match self {
                Self::NoPeer => "no connected peer".to_owned(),
                Self::Wire(err) => err.to_string(),
            },
        }
    }
}

/// Offers `message` to `send` until it is taken, while the handshake settles.
///
/// The window is the time the peer is given to appear, so it opens at the first refusal rather
/// than at the call: a send that is taken keeps the clock out of the publish path entirely.
async fn retry<Attempt: Offer>(
    mut attempt: Attempt,
    message: zeromq::ZmqMessage,
) -> Result<(), SendFailure> {
    let mut pending = message;
    let mut deadline = None;
    loop {
        match attempt.offer(pending).await {
            Ok(()) => return Ok(()),
            Err(WireError::ReturnToSender { message, .. }) => {
                match deadline {
                    None => deadline = Some(tokio::time::Instant::now() + SEND_RETRY_WINDOW),
                    Some(at) if tokio::time::Instant::now() >= at => {
                        return Err(SendFailure::NoPeer);
                    }
                    Some(_) => {}
                }
                pending = message;
                tokio::time::sleep(SEND_RETRY_STEP).await;
            }
            Err(err) => return Err(SendFailure::Wire(err)),
        }
    }
}

/// The send the client returns for a socket, borrowing the socket for `'socket`.
struct ClientSend;

impl<'socket> RefType<'socket> for ClientSend {
    type Ref = Pin<Box<dyn Future<Output = ZmqResult<()>> + Send + 'socket>>;
}

/// A socket at a fixed place on the heap, where the send in flight borrows it.
type Held<S> = Pin<AliasableBox<S>>;

/// A send in flight, stored beside the outbox it borrows its socket from.
type InFlightSend<S> = SelfieMut<'static, AliasableBox<S>, ClientSend>;

/// A sending socket whose send outlives the publish that started it.
///
/// The client's round-robin send takes the peer it picked out of rotation and puts it back only
/// once the frames are queued, so a send dropped while the peer applies back-pressure loses that
/// peer for good: with one peer, every later send reports no connected peer. Each send therefore
/// lives here rather than in the caller's future, borrowing the socket the outbox keeps on the
/// heap. A publish dropped mid-send leaves it parked, and the next publish completes it before
/// starting its own. The retry wait runs outside it, with the socket back here, so a publish
/// dropped while it waits sends nothing.
///
/// The socket is boxed once, when the outbox is made. A send is the client's own boxed future,
/// stored and polled as the client returned it, so a publish adds no allocation and no indirect
/// call to what the client already spends on it.
pub(crate) struct Outbox<S: 'static> {
    /// The socket while no send is in flight; the send holds it otherwise.
    socket: Option<Held<S>>,
    send: Option<InFlightSend<S>>,
    /// The name of a message whose publish was dropped mid-send, kept to report a failure of that
    /// send against it.
    dropped: Option<String>,
}

impl<S> Outbox<S>
where
    S: SocketSend + Send + Unpin + 'static,
{
    pub(crate) fn new(socket: S) -> Self {
        Self {
            socket: Some(AliasableBox::from_unique_pin(Box::pin(socket))),
            send: None,
            dropped: None,
        }
    }

    /// Hands the socket back, once a send a dropped publish left parked has completed.
    pub(crate) async fn into_socket(mut self) -> Option<S> {
        self.complete_dropped().await;
        self.socket
            .take()
            .map(|held| *AliasableBox::into_unique(Pin::into_inner(held)))
    }

    /// Sends `message` under `name` with the handshake retry, after completing a send whose
    /// publish was dropped.
    ///
    /// # Cancel safety
    ///
    /// Cancel-safe. Dropped before the socket takes the message, it sends nothing; dropped
    /// after, it leaves the send parked for the next call to complete.
    pub(crate) async fn send(
        &mut self,
        name: &str,
        message: zeromq::ZmqMessage,
    ) -> Result<(), ZmqError> {
        if self.send.is_some() {
            self.complete_dropped().await;
        }
        retry(Parked { outbox: self, name }, message)
            .await
            .map_err(|failure| failure.into_error(name))
    }

    /// Completes the send a dropped publish left parked and puts the socket back. Its outcome
    /// has no caller left, so a failure is logged against the message's name.
    async fn complete_dropped(&mut self) {
        let Some(send) = self.send.as_mut() else {
            return;
        };
        let outcome = poll_fn(|cx| send.with_referential_mut(|send| send.as_mut().poll(cx))).await;
        self.socket = self.send.take().map(SelfieMut::into_owned);
        let name = self.dropped.take().unwrap_or_default();
        if let Err(err) = outcome {
            let error = SendFailure::Wire(err).into_error(&name);
            tracing::warn!(
                name = %name,
                error = %error,
                "a message whose publish was dropped mid-send did not leave"
            );
        }
    }
}

/// A socket offered a message through its outbox, where a dropped offer stays parked.
struct Parked<'a, S: 'static> {
    outbox: &'a mut Outbox<S>,
    name: &'a str,
}

impl<S> Offer for Parked<'_, S>
where
    S: SocketSend + Send + Unpin + 'static,
{
    #[inline]
    fn offer(&mut self, message: zeromq::ZmqMessage) -> impl Future<Output = ZmqResult<()>> + Send {
        let socket = self
            .outbox
            .socket
            .take()
            .expect("a parked send completes before the next offer");
        self.outbox.send = Some(SelfieMut::new(socket, |socket| {
            Pin::into_inner(socket).send(message)
        }));
        Sending {
            outbox: &mut *self.outbox,
            name: self.name,
            done: false,
        }
    }
}

/// The send an offer started, polled in its outbox. Dropped before it completes, it leaves the
/// send parked there under the message's name.
struct Sending<'a, S: 'static> {
    outbox: &'a mut Outbox<S>,
    name: &'a str,
    done: bool,
}

impl<S> Future for Sending<'_, S> {
    type Output = ZmqResult<()>;

    #[inline]
    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<ZmqResult<()>> {
        let this = &mut *self;
        let send = this
            .outbox
            .send
            .as_mut()
            .expect("the send stays in the outbox until it completes");
        let outcome = ready!(send.with_referential_mut(|send| send.as_mut().poll(cx)));
        this.outbox.socket = this.outbox.send.take().map(SelfieMut::into_owned);
        this.done = true;
        Poll::Ready(outcome)
    }
}

impl<S> Drop for Sending<'_, S> {
    fn drop(&mut self) {
        if !self.done {
            self.outbox.dropped = Some(self.name.to_owned());
        }
    }
}

/// What keeps a subscription's deliveries coming, and stops them when the subscription goes.
///
/// Over sockets it is the driver task that owns the socket (every operation takes `&mut self`),
/// and dropping the handle aborts it, which is the only reliable teardown: a receive on a
/// peerless socket pends forever by design of the implementation. In process it is the
/// subscription's place on the transport, which dropping gives up.
///
/// The handle, not the task, holds the endpoint a bound subscription took: an aborted task is
/// dropped later, on whichever worker next runs it, so a slot held there would keep the endpoint
/// taken for a moment after the subscription is gone. Held here, it is free once the subscription
/// is dropped.
pub(crate) enum DriverHandle {
    Task(JoinHandle<()>, Option<BoundSlot>),
    /// The subscription's place on the transport, and the endpoint it holds when it binds, as
    /// its socket would hold the listener.
    #[cfg(feature = "testing")]
    InProcess(Registration, Option<BoundSlot>),
}

// A production subscription holds its driver task and the endpoint it bound, nothing beside them.
#[cfg(not(feature = "testing"))]
const _: () =
    assert!(size_of::<DriverHandle>() == size_of::<(JoinHandle<()>, Option<BoundSlot>)>());

impl Drop for DriverHandle {
    fn drop(&mut self) {
        match self {
            Self::Task(task, _slot) => task.abort(),
            #[cfg(feature = "testing")]
            Self::InProcess(registration, _slot) => registration.give_up(),
        }
    }
}

/// The driver task's end of a subscription's delivery channel.
pub(crate) type DeliverySender = mpsc::Sender<Result<ZmqMessage, ZmqError>>;

/// The handler's end of a subscription's delivery channel: the driver task's bounded channel, or,
/// in process, the channel the transport hands deliveries to.
///
/// Without the `testing` feature it is the driver's receiver itself: one variant, no tag, no
/// branch.
pub(crate) enum DeliveryReceiver {
    Driver(mpsc::Receiver<Result<ZmqMessage, ZmqError>>),
    #[cfg(feature = "testing")]
    InProcess(Deliveries),
}

// A production subscription reads its driver's channel and nothing beside it.
#[cfg(not(feature = "testing"))]
const _: () = assert!(
    size_of::<DeliveryReceiver>() == size_of::<mpsc::Receiver<Result<ZmqMessage, ZmqError>>>()
);

impl DeliveryReceiver {
    fn poll_recv(&mut self, cx: &mut Context<'_>) -> Poll<Option<Result<ZmqMessage, ZmqError>>> {
        match self {
            Self::Driver(rx) => rx.poll_recv(cx),
            #[cfg(feature = "testing")]
            Self::InProcess(rx) => rx.poll_recv(cx),
        }
    }
}

/// A sending socket a publisher attached, or, in process, the transport it sends over and
/// whether its sends reach a subscription of this service (the loopback arrangement) or a peer
/// outside it.
///
/// Without the `testing` feature it is the socket itself: one variant, no tag, no branch.
pub(crate) enum Sender<S> {
    Socket(S),
    #[cfg(feature = "testing")]
    InProcess {
        bus: Arc<Bus>,
        local: bool,
    },
}

/// The socket side of any subscription: the driver task's channel, one delivery at a time, which
/// is all a receive on a ZMTP socket yields.
///
/// The public subscribers wrap it - the batching patterns through the framework's client-side
/// buffer, the request-reply one directly.
pub(crate) struct WireSubscriber {
    pub(crate) rx: DeliveryReceiver,
    pub(crate) _driver: DriverHandle,
}

impl Subscriber for WireSubscriber {
    type Message = ZmqMessage;
    type Error = ZmqError;

    fn stream(&mut self) -> impl Stream<Item = Result<ZmqMessage, ZmqError>> + Send + '_ {
        // Poll the channel in place rather than wrapping it in an owning stream, so `stream`
        // can be called again after the returned stream is dropped (the runtime and the
        // conformance helpers re-enter it per call).
        futures::stream::poll_fn(move |cx| self.rx.poll_recv(cx))
    }
}

pub(crate) type SharedLifecycle = Arc<Lifecycle>;

#[cfg(test)]
mod tests {
    use std::pin::pin;
    use std::sync::Mutex;

    use async_trait::async_trait;
    use bytes::Bytes;
    use futures::poll;
    use tokio::sync::Notify;
    use tokio::time::advance;

    use super::*;

    /// A socket that refuses the first send the way a peerless one does, after taking `stall` off
    /// the clock: the handshake a freshly dialled socket is still settling takes time of its own.
    struct Handshaking {
        refusals: usize,
        stall: Duration,
        stalled: bool,
    }

    #[async_trait]
    impl SocketSend for Handshaking {
        async fn send(&mut self, message: zeromq::ZmqMessage) -> ZmqResult<()> {
            if self.refusals == 0 {
                return Ok(());
            }
            if !self.stalled {
                self.stalled = true;
                advance(self.stall).await;
            }
            self.refusals -= 1;
            Err(WireError::ReturnToSender {
                reason: "no peer",
                message,
            })
        }
    }

    /// The window is the time the peer is given to appear, so it starts when the peer first
    /// refuses. A send that took longer than the window to reach that refusal has not used the
    /// window up - and a send that succeeds asks the clock nothing at all.
    #[tokio::test(start_paused = true)]
    async fn the_retry_window_starts_when_the_peer_first_refuses() {
        let mut socket = Handshaking {
            refusals: 1,
            stall: SEND_RETRY_WINDOW + Duration::from_secs(1),
            stalled: false,
        };

        send_with_retry(
            &mut socket,
            "orders",
            zeromq::ZmqMessage::from(Bytes::from_static(b"{}")),
        )
        .await
        .expect("the peer refused once and then took the message");
    }

    /// A socket whose first send waits for the peer to open a gate, as one under back-pressure
    /// does, and which records the payload of every send it completes.
    struct Gated {
        gate: Arc<Notify>,
        held: bool,
        sent: Arc<Mutex<Vec<Bytes>>>,
    }

    #[async_trait]
    impl SocketSend for Gated {
        async fn send(&mut self, message: zeromq::ZmqMessage) -> ZmqResult<()> {
            if self.held {
                self.gate.notified().await;
                self.held = false;
            }
            let payload = message.get(0).cloned().unwrap_or_default();
            self.sent.lock().expect("not poisoned").push(payload);
            Ok(())
        }
    }

    /// A send dropped while the socket holds it stays with the outbox, and the next send
    /// completes it before its own: nothing the dropped publish handed over is lost or reordered.
    #[tokio::test]
    async fn a_send_dropped_mid_flight_completes_before_the_next() {
        let gate = Arc::new(Notify::new());
        let sent = Arc::new(Mutex::new(Vec::new()));
        let mut outbox = Outbox::new(Gated {
            gate: Arc::clone(&gate),
            held: true,
            sent: Arc::clone(&sent),
        });

        {
            let dropped = pin!(outbox.send(
                "first",
                zeromq::ZmqMessage::from(Bytes::from_static(b"first"))
            ));
            assert!(
                poll!(dropped).is_pending(),
                "the socket holds the first send"
            );
        }
        gate.notify_one();
        outbox
            .send(
                "second",
                zeromq::ZmqMessage::from(Bytes::from_static(b"second")),
            )
            .await
            .expect("the second send leaves");

        assert_eq!(
            *sent.lock().expect("not poisoned"),
            [Bytes::from_static(b"first"), Bytes::from_static(b"second")],
        );
    }
}