44
55//! Forward UDS connections to a remote yamux endpoint over TCP.
66//!
7- //! Maintains a single TCP connection and multiplexes each incoming UDS
7+ //! Maintains a pool of TCP connections and multiplexes each incoming UDS
88//! connection onto a yamux stream.
99
1010use anyhow:: { Context , Result } ;
1111use clap:: Parser ;
12+ use std:: sync:: atomic:: { AtomicUsize , Ordering } ;
13+ use std:: sync:: Arc ;
1214use tokio:: net:: { TcpStream , UnixListener } ;
1315use tokio:: sync:: { mpsc, oneshot} ;
1416use tokio_util:: compat:: { FuturesAsyncReadCompatExt , TokioAsyncReadCompatExt } ;
@@ -24,8 +26,14 @@ struct Args {
2426 /// Yamux server address (host:port)
2527#[ arg( long) ]
2628 yamux_addr : String ,
29+
30+ /// Number of TCP connections in the yamux pool
31+ #[ arg( long, default_value = "1" ) ]
32+ yamux_conns : usize ,
2733}
2834
35+ type YamuxRequest = oneshot:: Sender < anyhow:: Result < yamux:: Stream > > ;
36+
2937async fn connect_yamux (
3038 addr : & str ,
3139) -> Result < yamux:: Connection < tokio_util:: compat:: Compat < TcpStream > > > {
@@ -44,60 +52,135 @@ async fn connect_yamux(
4452 ) )
4553}
4654
47- async fn drive_yamux (
48- mut conn : yamux:: Connection < tokio_util:: compat:: Compat < TcpStream > > ,
49- mut requests : mpsc:: Receiver < oneshot:: Sender < anyhow:: Result < yamux:: Stream > > > ,
50- ) {
55+ async fn drive_yamux ( addr : String , mut requests : mpsc:: Receiver < YamuxRequest > ) {
56+ use futures:: FutureExt ;
57+
58+ let mut pending: Vec < YamuxRequest > = Vec :: new ( ) ;
59+ let mut backoff = std:: time:: Duration :: from_secs ( 1 ) ;
60+ let max_backoff = std:: time:: Duration :: from_secs ( 60 ) ;
61+
5162 loop {
52- tokio:: select! {
53- inbound = std:: future:: poll_fn( |cx| conn. poll_next_inbound( cx) ) => {
54- match inbound {
55- Some ( Ok ( stream) ) => {
56- info!( "yamux: inbound stream {stream}" ) ;
57- }
58- Some ( Err ( e) ) => {
59- error!( "yamux: inbound error: {e}" ) ;
60- break ;
61- }
62- None => {
63- info!( "yamux: connection closed by peer" ) ;
64- break ;
63+ let mut conn = loop {
64+ match connect_yamux ( & addr) . await {
65+ Ok ( conn) => {
66+ backoff = std:: time:: Duration :: from_secs ( 1 ) ;
67+ info ! ( "yamux: connected to {addr}" ) ;
68+ break conn;
69+ }
70+ Err ( e) => {
71+ error ! ( "yamux: connect error: {e}" ) ;
72+ let delay = tokio:: time:: sleep ( backoff) ;
73+ tokio:: pin!( delay) ;
74+ tokio:: select! {
75+ _ = & mut delay => {
76+ backoff = std:: cmp:: min( backoff * 2 , max_backoff) ;
77+ }
78+ request = requests. recv( ) => {
79+ match request {
80+ Some ( reply) => pending. push( reply) ,
81+ None => return ,
82+ }
83+ }
6584 }
6685 }
6786 }
68- request = requests. recv( ) => {
69- let Some ( reply) = request else {
70- break ;
71- } ;
72- let result = std:: future:: poll_fn( |cx| conn. poll_new_outbound( cx) )
73- . await
74- . map_err( |e| anyhow:: anyhow!( "failed to open yamux stream: {e}" ) ) ;
75- let _ = reply. send( result) ;
87+ } ;
88+
89+ loop {
90+ let request_fut = if let Some ( reply) = pending. pop ( ) {
91+ futures:: future:: ready ( Some ( reply) ) . boxed ( )
92+ } else {
93+ requests. recv ( ) . boxed ( )
94+ } ;
95+
96+ tokio:: select! {
97+ inbound = std:: future:: poll_fn( |cx| conn. poll_next_inbound( cx) ) => {
98+ match inbound {
99+ Some ( Ok ( stream) ) => {
100+ info!( "yamux: inbound stream {stream}" ) ;
101+ }
102+ Some ( Err ( e) ) => {
103+ error!( "yamux: inbound error: {e}" ) ;
104+ break ;
105+ }
106+ None => {
107+ info!( "yamux: connection closed by peer" ) ;
108+ break ;
109+ }
110+ }
111+ }
112+ request = request_fut => {
113+ let Some ( reply) = request else {
114+ return ;
115+ } ;
116+ let result = std:: future:: poll_fn( |cx| conn. poll_new_outbound( cx) )
117+ . await
118+ . map_err( |e| anyhow:: anyhow!( "failed to open yamux stream: {e}" ) ) ;
119+ let _ = reply. send( result) ;
120+ }
76121 }
77122 }
78123 }
79124}
80125
126+ async fn spawn_yamux_driver ( addr : & str ) -> Result < mpsc:: Sender < YamuxRequest > > {
127+ let ( request_tx, request_rx) = mpsc:: channel ( 128 ) ;
128+ tokio:: spawn ( drive_yamux ( addr. to_string ( ) , request_rx) ) ;
129+ Ok ( request_tx)
130+ }
131+
132+ struct YamuxPool {
133+ senders : Vec < mpsc:: Sender < YamuxRequest > > ,
134+ next : AtomicUsize ,
135+ }
136+
137+ impl YamuxPool {
138+ async fn new ( addr : & str , num_conns : usize ) -> Result < Self > {
139+ let mut senders = Vec :: with_capacity ( num_conns) ;
140+ for _ in 0 ..num_conns {
141+ senders. push ( spawn_yamux_driver ( addr) . await ?) ;
142+ }
143+ Ok ( Self {
144+ senders,
145+ next : AtomicUsize :: new ( 0 ) ,
146+ } )
147+ }
148+
149+ async fn open_stream ( & self ) -> Result < yamux:: Stream > {
150+ let idx = self . next . fetch_add ( 1 , Ordering :: Relaxed ) % self . senders . len ( ) ;
151+ let ( reply_tx, reply_rx) = oneshot:: channel ( ) ;
152+ self . senders [ idx]
153+ . send ( reply_tx)
154+ . await
155+ . map_err ( |_| anyhow:: anyhow!( "yamux driver closed" ) ) ?;
156+ reply_rx
157+ . await
158+ . map_err ( |_| anyhow:: anyhow!( "yamux driver dropped response channel" ) ) ?
159+ }
160+ }
161+
81162#[ tokio:: main]
82163async fn main ( ) -> Result < ( ) > {
83164 tracing_subscriber:: fmt:: init ( ) ;
84165 let args = Args :: parse ( ) ;
166+ if args. yamux_conns == 0 {
167+ anyhow:: bail!( "yamux_conns must be >= 1" ) ;
168+ }
85169
86170 let _ = std:: fs:: remove_file ( & args. uds ) ;
87171
88172 let listener = UnixListener :: bind ( & args. uds )
89173 . with_context ( || format ! ( "failed to bind UDS at {}" , args. uds) ) ?;
90174
91- let connection = connect_yamux ( & args. yamux_addr )
92- . await
93- . context ( "failed to connect yamux" ) ?;
94-
95- let ( request_tx, request_rx) = mpsc:: channel ( 128 ) ;
96- tokio:: spawn ( drive_yamux ( connection, request_rx) ) ;
175+ let pool = Arc :: new (
176+ YamuxPool :: new ( & args. yamux_addr , args. yamux_conns )
177+ . await
178+ . context ( "failed to connect yamux" ) ?,
179+ ) ;
97180
98181 info ! (
99- "listening on {}, forwarding to yamux {}" ,
100- args. uds, args. yamux_addr
182+ "listening on {}, forwarding to yamux {} (pool size {}) " ,
183+ args. uds, args. yamux_addr, args . yamux_conns
101184 ) ;
102185
103186 loop {
@@ -106,24 +189,14 @@ async fn main() -> Result<()> {
106189 . await
107190 . context ( "failed to accept UDS connection" ) ?;
108191
109- let request_tx = request_tx . clone ( ) ;
192+ let pool = pool . clone ( ) ;
110193 tokio:: spawn ( async move {
111- let ( reply_tx, reply_rx) = oneshot:: channel ( ) ;
112- if request_tx. send ( reply_tx) . await . is_err ( ) {
113- error ! ( "yamux driver closed" ) ;
114- return ;
115- }
116-
117- let stream = match reply_rx. await {
118- Ok ( Ok ( stream) ) => stream,
119- Ok ( Err ( e) ) => {
194+ let stream = match pool. open_stream ( ) . await {
195+ Ok ( stream) => stream,
196+ Err ( e) => {
120197 error ! ( "failed to open yamux stream: {e}" ) ;
121198 return ;
122199 }
123- Err ( _) => {
124- error ! ( "yamux driver dropped response channel" ) ;
125- return ;
126- }
127200 } ;
128201
129202 let stream = stream. compat ( ) ;
0 commit comments