├── rustfmt.toml ├── .gitignore ├── examples ├── printreq.py ├── reqhandler.py ├── sendresp.py ├── get.py ├── basichandler.py ├── getstream.py ├── streamhandler.py └── holdhandler.py ├── CHANGELOG.md ├── Cargo.toml ├── src ├── net.rs ├── shuffle.rs ├── pool.rs ├── lib.rs ├── track.rs ├── waker.rs ├── listener.rs ├── list.rs ├── resolver.rs ├── app.rs ├── zmq.rs ├── executor.rs ├── timer.rs ├── arena.rs └── main.rs ├── benches ├── server.rs └── client.rs ├── README.md └── COPYING /rustfmt.toml: -------------------------------------------------------------------------------- 1 | # using the defaults 2 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | 2 | /target 3 | **/*.rs.bk 4 | -------------------------------------------------------------------------------- /examples/printreq.py: -------------------------------------------------------------------------------- 1 | # this handler just outputs the request ID 2 | 3 | import tnetstring 4 | import zmq 5 | 6 | ctx = zmq.Context() 7 | sock = ctx.socket(zmq.PULL) 8 | sock.connect('ipc://client-out') 9 | 10 | while True: 11 | m = sock.recv_multipart() 12 | req = tnetstring.loads(m[0][1:]) 13 | print('{} {}'.format(req[b'from'].decode(), req[b'id'].decode())) 14 | -------------------------------------------------------------------------------- /examples/reqhandler.py: -------------------------------------------------------------------------------- 1 | # this handler responds to every request with "hello world" 2 | 3 | import tnetstring 4 | import zmq 5 | 6 | ctx = zmq.Context() 7 | sock = ctx.socket(zmq.REP) 8 | sock.connect('ipc://client') 9 | 10 | while True: 11 | m_raw = sock.recv() 12 | req = tnetstring.loads(m_raw[1:]) 13 | print('IN {}'.format(req)) 14 | 15 | resp = {} 16 | resp[b'id'] = req[b'id'] 17 | resp[b'code'] = 200 18 | resp[b'reason'] = b'OK' 19 | resp[b'headers'] = [[b'Content-Type', b'text/plain']] 20 | resp[b'body'] = b'hello world\n' 21 | 22 | print('OUT {}'.format(resp)) 23 | sock.send(b'T' + tnetstring.dumps(resp)) 24 | -------------------------------------------------------------------------------- /examples/sendresp.py: -------------------------------------------------------------------------------- 1 | # this program sends a response to a certain request ID 2 | 3 | import sys 4 | import time 5 | import tnetstring 6 | import zmq 7 | 8 | body = sys.argv[1] 9 | addr = sys.argv[2].encode() 10 | rid = sys.argv[3].encode() 11 | 12 | ctx = zmq.Context() 13 | sock = ctx.socket(zmq.PUB) 14 | sock.connect('ipc://client-in') 15 | 16 | # await subscription 17 | time.sleep(0.01) 18 | 19 | resp = {} 20 | resp[b'from'] = b'sendresp' 21 | resp[b'id'] = rid 22 | resp[b'code'] = 200 23 | resp[b'reason'] = b'OK' 24 | resp[b'headers'] = [[b'Content-Type', b'text/plain']] 25 | resp[b'body'] = '{}\n'.format(body).encode() 26 | 27 | m = [addr + b' T' + tnetstring.dumps(resp)] 28 | 29 | sock.send_multipart(m) 30 | -------------------------------------------------------------------------------- /examples/get.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import uuid 3 | import tnetstring 4 | import zmq 5 | 6 | if len(sys.argv) < 2: 7 | print('usage: {} [url]'.format(sys.argv[0])) 8 | sys.exit(1) 9 | 10 | ctx = zmq.Context() 11 | sock = ctx.socket(zmq.REQ) 12 | sock.connect('ipc://server') 13 | 14 | req = { 15 | b'method': b'GET', 16 | b'uri': sys.argv[1].encode('utf-8'), 17 | #b'follow-redirects': True, 18 | #b'ignore-tls-errors': True, 19 | } 20 | 21 | sock.send(b'T' + tnetstring.dumps(req)) 22 | 23 | resp = tnetstring.loads(sock.recv()[1:]) 24 | if b'type' in resp and resp[b'type'] == b'error': 25 | print('error: {}'.format(resp[b'condition'])) 26 | sys.exit(1) 27 | 28 | print('code={} reason=[{}]'.format(resp[b'code'], resp[b'reason'])) 29 | for h in resp[b'headers']: 30 | print('{}: {}'.format(h[0], h[1])) 31 | 32 | if b'body' in resp: 33 | print('\n{}'.format(resp[b'body'])) 34 | else: 35 | print('\n') 36 | -------------------------------------------------------------------------------- /examples/basichandler.py: -------------------------------------------------------------------------------- 1 | # this handler responds to every request with "hello world" 2 | 3 | import os 4 | import time 5 | import tnetstring 6 | import zmq 7 | 8 | instance_id = 'basichandler.{}'.format(os.getpid()).encode() 9 | 10 | ctx = zmq.Context() 11 | in_sock = ctx.socket(zmq.PULL) 12 | in_sock.connect('ipc://client-out') 13 | out_sock = ctx.socket(zmq.PUB) 14 | out_sock.connect('ipc://client-in') 15 | 16 | # await subscription 17 | time.sleep(0.01) 18 | 19 | while True: 20 | m_raw = in_sock.recv() 21 | req = tnetstring.loads(m_raw[1:]) 22 | print('IN {}'.format(req)) 23 | 24 | resp = {} 25 | resp[b'from'] = instance_id 26 | resp[b'id'] = req[b'id'] 27 | resp[b'code'] = 200 28 | resp[b'reason'] = b'OK' 29 | resp[b'headers'] = [[b'Content-Type', b'text/plain']] 30 | resp[b'body'] = b'hello world\n' 31 | 32 | print('OUT {}'.format(resp)) 33 | out_sock.send(req[b'from'] + b' T' + tnetstring.dumps(resp)) 34 | -------------------------------------------------------------------------------- /examples/getstream.py: -------------------------------------------------------------------------------- 1 | import sys 2 | import time 3 | import uuid 4 | import tnetstring 5 | import zmq 6 | 7 | client_id = b'getstream.py' 8 | 9 | ctx = zmq.Context() 10 | out_sock = ctx.socket(zmq.PUSH) 11 | out_sock.connect('ipc://server-in') 12 | out_stream_sock = ctx.socket(zmq.ROUTER) 13 | out_stream_sock.connect('ipc://server-in-stream') 14 | in_sock = ctx.socket(zmq.SUB) 15 | in_sock.setsockopt(zmq.SUBSCRIBE, client_id) 16 | in_sock.connect('ipc://server-out') 17 | 18 | time.sleep(0.5) 19 | 20 | rid = str(uuid.uuid4()).encode('utf-8') 21 | inseq = 0 22 | outseq = 0 23 | out_sock.send(b'T' + tnetstring.dumps({ 24 | b'from': client_id, 25 | b'id': rid, 26 | b'seq': outseq, 27 | b'method': b'GET', 28 | b'uri': sys.argv[1].encode('utf-8'), 29 | b'stream': True, 30 | b'credits': 8192, 31 | })) 32 | outseq += 1 33 | 34 | while True: 35 | buf = in_sock.recv() 36 | at = buf.find(b' ') 37 | receiver = buf[:at] 38 | indata = tnetstring.loads(buf[at + 2:]) 39 | if indata[b'id'] != rid: 40 | continue 41 | print('IN: {}'.format(indata)) 42 | assert(indata[b'seq'] == inseq) 43 | inseq += 1 44 | if (b'type' in indata and (indata[b'type'] == b'error' or indata[b'type'] == b'cancel')) or (b'type' not in indata and b'more' not in indata): 45 | break 46 | raddr = indata[b'from'] 47 | if b'body' in indata and len(indata[b'body']) > 0: 48 | outdata = { 49 | b'id': rid, 50 | b'from': client_id, 51 | b'seq': outseq, 52 | b'type': b'credit', 53 | b'credits': len(indata[b'body']), 54 | } 55 | print('OUT: {}'.format(outdata)) 56 | out_stream_sock.send_multipart([raddr, b'', b'T' + tnetstring.dumps(outdata)]) 57 | outseq += 1 58 | -------------------------------------------------------------------------------- /CHANGELOG.md: -------------------------------------------------------------------------------- 1 | Condure Changelog 2 | ================= 3 | 4 | v. 1.10.1 (2024-08-12) 5 | 6 | * Fix build with recent rustc. 7 | 8 | v. 1.10.0 (2023-06-29) 9 | 10 | * Add support for outgoing connections. 11 | * Ability to set mode/user/group when listening on Unix socket. 12 | 13 | v. 1.9.2 (2023-02-06) 14 | 15 | * Fix WebSocket compression with fragmented messages. 16 | 17 | v. 1.9.1 (2023-01-19) 18 | 19 | * Fix crash in stream connection handler. 20 | 21 | v. 1.9.0 (2022-12-05) 22 | 23 | * Support permessage-deflate WebSocket compression. 24 | 25 | v. 1.8.0 (2022-11-03) 26 | 27 | * Compatibility with httparse 1.8. 28 | * Add more benchmarks. 29 | 30 | v. 1.7.0 (2022-08-18) 31 | 32 | * Fix worker thread hang when backend buffer is full. 33 | 34 | v. 1.6.0 (2022-04-15) 35 | 36 | * Significantly reduce connection memory usage. 37 | * Allow up to 64 headers in requests and responses. 38 | * Allow WebSocket requests that include a Content-Length of 0. 39 | 40 | v. 1.5.0 (2022-03-11) 41 | 42 | * Ability to listen on a Unix socket for client connections. 43 | 44 | v. 1.4.1 (2021-10-24) 45 | 46 | * Fix crash when sending too fast to clients. 47 | 48 | v. 1.4.0 (2021-10-22) 49 | 50 | * Port connection handler to use async functions. 51 | 52 | v. 1.3.1 (2021-08-11) 53 | 54 | * Fixes for high load. 55 | 56 | v. 1.3.0 (2021-07-29) 57 | 58 | * Port to async/await. 59 | 60 | v. 1.2.0 (2021-05-04) 61 | 62 | * Send PING/PONG frame data to clients. 63 | * Port to mio 0.7. 64 | 65 | v. 1.1.0 (2020-11-02) 66 | 67 | * TLS support. 68 | * Don't preallocate connection buffers. 69 | * Start using async/await in some places. 70 | 71 | v. 1.0.1 (2020-07-24) 72 | 73 | * Remove some unsafe usage. 74 | 75 | v. 1.0.0 (2020-07-21) 76 | 77 | * Stable version. 78 | -------------------------------------------------------------------------------- /Cargo.toml: -------------------------------------------------------------------------------- 1 | [package] 2 | name = "condure" 3 | version = "1.10.1" 4 | authors = ["Justin Karneges "] 5 | license = "Apache-2.0" 6 | description = "HTTP/WebSocket connection manager" 7 | repository = "https://github.com/fanout/condure" 8 | readme = "README.md" 9 | edition = "2018" 10 | autobins = false 11 | 12 | [[bin]] 13 | name = "condure" 14 | 15 | [profile.dev] 16 | panic = "abort" 17 | 18 | [profile.release] 19 | panic = "abort" 20 | 21 | [dependencies] 22 | arrayvec = "0.7" 23 | base64 = "0.13" 24 | clap = { version = "4.3", features = ["cargo", "string", "wrap_help"] } 25 | httparse = "1.7" 26 | ipnet = "2" 27 | libc = "0.2" 28 | log = "0.4" 29 | miniz_oxide = "0.6" 30 | mio = { version = "0.8", features = ["os-poll", "os-ext", "net"] } 31 | openssl = "0.10" 32 | paste = "1.0" 33 | sha1 = "0.10" 34 | signal-hook = "0.3" 35 | slab = "0.4" 36 | socket2 = "0.4" 37 | thiserror = "1.0" 38 | time = { version = "0.3", features = ["formatting", "local-offset", "macros"] } 39 | url = "2.3" 40 | zmq = "0.9" 41 | 42 | [dev-dependencies] 43 | criterion = "0.5" 44 | env_logger = { version = "0.9", default-features = false } 45 | test-log = "0.2" 46 | 47 | [[bench]] 48 | name = "server" 49 | harness = false 50 | 51 | [[bench]] 52 | name = "client" 53 | harness = false 54 | 55 | [package.metadata.deb] 56 | extended-description = """\ 57 | Condure is a service that manages network connections on behalf of server 58 | applications, in order to allow controlling the connections from multiple 59 | processes. Applications communicate with Condure over ZeroMQ. 60 | 61 | Condure can only manage connections for protocols it knows about. Currently 62 | this is HTTP/1 and WebSockets. 63 | 64 | The project was inspired by Mongrel2. 65 | """ 66 | separate-debug-symbols = true 67 | build-depends = "pkg-config, libzmq3-dev" 68 | 69 | [package.metadata.rpm] 70 | package = "condure" 71 | 72 | [package.metadata.rpm.cargo] 73 | buildflags = ["--release"] 74 | 75 | [package.metadata.rpm.targets] 76 | condure = { path = "/usr/bin/condure" } 77 | -------------------------------------------------------------------------------- /src/net.rs: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2022 Fanout, Inc. 3 | * 4 | * Licensed under the Apache License, Version 2.0 (the "License"); 5 | * you may not use this file except in compliance with the License. 6 | * You may obtain a copy of the License at 7 | * 8 | * http://www.apache.org/licenses/LICENSE-2.0 9 | * 10 | * Unless required by applicable law or agreed to in writing, software 11 | * distributed under the License is distributed on an "AS IS" BASIS, 12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | * See the License for the specific language governing permissions and 14 | * limitations under the License. 15 | */ 16 | 17 | use log::error; 18 | use mio::net::{TcpListener, TcpStream, UnixListener, UnixStream}; 19 | use socket2::Socket; 20 | use std::fmt; 21 | use std::os::unix::io::{FromRawFd, IntoRawFd}; 22 | use std::ptr; 23 | 24 | pub fn set_socket_opts(stream: &mut TcpStream) { 25 | if let Err(e) = stream.set_nodelay(true) { 26 | error!("set nodelay failed: {:?}", e); 27 | } 28 | 29 | // safety: we move the value out of stream and replace it at the end 30 | let ret = unsafe { 31 | let s = ptr::read(stream); 32 | let socket = Socket::from_raw_fd(s.into_raw_fd()); 33 | let ret = socket.set_keepalive(true); 34 | ptr::write(stream, TcpStream::from_raw_fd(socket.into_raw_fd())); 35 | 36 | ret 37 | }; 38 | 39 | if let Err(e) = ret { 40 | error!("set keepalive failed: {:?}", e); 41 | } 42 | } 43 | 44 | #[derive(Debug)] 45 | pub enum SocketAddr { 46 | Ip(std::net::SocketAddr), 47 | Unix(mio::net::SocketAddr), 48 | } 49 | 50 | impl fmt::Display for SocketAddr { 51 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 52 | match self { 53 | Self::Ip(a) => write!(f, "{}", a), 54 | Self::Unix(a) => write!(f, "{:?}", a), 55 | } 56 | } 57 | } 58 | 59 | #[derive(Debug)] 60 | pub enum NetListener { 61 | Tcp(TcpListener), 62 | Unix(UnixListener), 63 | } 64 | 65 | #[derive(Debug)] 66 | pub enum NetStream { 67 | Tcp(TcpStream), 68 | Unix(UnixStream), 69 | } 70 | -------------------------------------------------------------------------------- /src/shuffle.rs: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (c) 2016 Alex Crichton 3 | * Copyright (c) 2017 The Tokio Authors 4 | * 5 | * Licensed under the Apache License, Version 2.0 (the "License"); 6 | * you may not use this file except in compliance with the License. 7 | * You may obtain a copy of the License at 8 | * 9 | * http://www.apache.org/licenses/LICENSE-2.0 10 | * 11 | * Unless required by applicable law or agreed to in writing, software 12 | * distributed under the License is distributed on an "AS IS" BASIS, 13 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 14 | * See the License for the specific language governing permissions and 15 | * limitations under the License. 16 | */ 17 | 18 | use std::{ 19 | cell::Cell, 20 | collections::hash_map::DefaultHasher, 21 | hash::Hasher, 22 | num::Wrapping, 23 | sync::atomic::{AtomicUsize, Ordering}, 24 | }; 25 | 26 | // Based on [Fisher–Yates shuffle]. 27 | // 28 | // [Fisher–Yates shuffle]: https://en.wikipedia.org/wiki/Fisher–Yates_shuffle 29 | pub fn shuffle(slice: &mut [T]) { 30 | for i in (1..slice.len()).rev() { 31 | slice.swap(i, gen_index(i + 1)); 32 | } 33 | } 34 | 35 | /// Return a value from `0..n`. 36 | fn gen_index(n: usize) -> usize { 37 | (random() % n as u64) as usize 38 | } 39 | 40 | /// Pseudorandom number generator based on [xorshift*]. 41 | /// 42 | /// [xorshift*]: https://en.wikipedia.org/wiki/Xorshift#xorshift* 43 | pub fn random() -> u64 { 44 | thread_local! { 45 | static RNG: Cell> = Cell::new(Wrapping(prng_seed())); 46 | } 47 | 48 | fn prng_seed() -> u64 { 49 | static COUNTER: AtomicUsize = AtomicUsize::new(0); 50 | 51 | // Any non-zero seed will do 52 | let mut seed = 0; 53 | while seed == 0 { 54 | let mut hasher = DefaultHasher::new(); 55 | hasher.write_usize(COUNTER.fetch_add(1, Ordering::Relaxed)); 56 | seed = hasher.finish(); 57 | } 58 | seed 59 | } 60 | 61 | RNG.with(|rng| { 62 | let mut x = rng.get(); 63 | debug_assert_ne!(x.0, 0); 64 | x ^= x >> 12; 65 | x ^= x << 25; 66 | x ^= x >> 27; 67 | rng.set(x); 68 | x.0.wrapping_mul(0x2545_f491_4f6c_dd1d) 69 | }) 70 | } 71 | -------------------------------------------------------------------------------- /examples/streamhandler.py: -------------------------------------------------------------------------------- 1 | # this handler responds to every request with "hello world" 2 | 3 | import os 4 | import tnetstring 5 | import zmq 6 | 7 | instance_id = 'streamhandler.{}'.format(os.getpid()).encode('utf-8') 8 | 9 | ctx = zmq.Context() 10 | in_sock = ctx.socket(zmq.PULL) 11 | in_sock.connect('ipc://client-out') 12 | in_stream_sock = ctx.socket(zmq.ROUTER) 13 | in_stream_sock.identity = instance_id 14 | in_stream_sock.connect('ipc://client-out-stream') 15 | out_sock = ctx.socket(zmq.PUB) 16 | out_sock.connect('ipc://client-in') 17 | 18 | poller = zmq.Poller() 19 | poller.register(in_sock, zmq.POLLIN) 20 | poller.register(in_stream_sock, zmq.POLLIN) 21 | 22 | while True: 23 | socks = dict(poller.poll(None)) 24 | 25 | if socks.get(in_sock) == zmq.POLLIN: 26 | m_raw = in_sock.recv() 27 | elif socks.get(in_stream_sock) == zmq.POLLIN: 28 | m_list = in_stream_sock.recv_multipart() 29 | m_raw = m_list[2] 30 | else: 31 | continue 32 | 33 | req = tnetstring.loads(m_raw[1:]) 34 | print('IN {}'.format(req)) 35 | 36 | if req.get(b'type'): 37 | # skip all non-data messages 38 | continue 39 | 40 | if req.get(b'uri', b'').startswith(b'ws'): 41 | resp = {} 42 | resp[b'from'] = instance_id 43 | resp[b'id'] = req[b'id'] 44 | resp[b'seq'] = 0 45 | resp[b'code'] = 101 46 | resp[b'reason'] = b'Switching Protocols' 47 | resp[b'credits'] = 1024 48 | 49 | print('OUT {} {}'.format(req[b'from'], resp)) 50 | out_sock.send(req[b'from'] + b' T' + tnetstring.dumps(resp)) 51 | 52 | resp = {} 53 | resp[b'from'] = instance_id 54 | resp[b'id'] = req[b'id'] 55 | resp[b'seq'] = 1 56 | resp[b'body'] = b'hello world' 57 | 58 | print('OUT {} {}'.format(req[b'from'], resp)) 59 | out_sock.send(req[b'from'] + b' T' + tnetstring.dumps(resp)) 60 | 61 | resp = {} 62 | resp[b'from'] = instance_id 63 | resp[b'id'] = req[b'id'] 64 | resp[b'seq'] = 2 65 | resp[b'type'] = b'close' 66 | 67 | print('OUT {} {}'.format(req[b'from'], resp)) 68 | out_sock.send(req[b'from'] + b' T' + tnetstring.dumps(resp)) 69 | else: 70 | resp = {} 71 | resp[b'from'] = instance_id 72 | resp[b'id'] = req[b'id'] 73 | resp[b'seq'] = 0 74 | resp[b'code'] = 200 75 | resp[b'reason'] = b'OK' 76 | resp[b'headers'] = [[b'Content-Type', b'text/plain']] 77 | resp[b'more'] = True 78 | resp[b'credits'] = 1024 79 | 80 | print('OUT {} {}'.format(req[b'from'], resp)) 81 | out_sock.send(req[b'from'] + b' T' + tnetstring.dumps(resp)) 82 | 83 | resp = {} 84 | resp[b'from'] = instance_id 85 | resp[b'id'] = req[b'id'] 86 | resp[b'seq'] = 1 87 | resp[b'body'] = b'hello world\n' 88 | 89 | print('OUT {} {}'.format(req[b'from'], resp)) 90 | out_sock.send(req[b'from'] + b' T' + tnetstring.dumps(resp)) 91 | -------------------------------------------------------------------------------- /examples/holdhandler.py: -------------------------------------------------------------------------------- 1 | # this handler holds all connections open 2 | 3 | import os 4 | import time 5 | import datetime 6 | import calendar 7 | import tnetstring 8 | import zmq 9 | 10 | CONN_TTL = 60000 11 | EXPIRE_INTERVAL = 60000 12 | 13 | instance_id = 'holdhandler.{}'.format(os.getpid()).encode('utf-8') 14 | 15 | ctx = zmq.Context() 16 | in_sock = ctx.socket(zmq.PULL) 17 | in_sock.connect('ipc://client-out') 18 | in_stream_sock = ctx.socket(zmq.ROUTER) 19 | in_stream_sock.identity = instance_id 20 | in_stream_sock.connect('ipc://client-out-stream') 21 | out_sock = ctx.socket(zmq.PUB) 22 | out_sock.connect('ipc://client-in') 23 | 24 | poller = zmq.Poller() 25 | poller.register(in_sock, zmq.POLLIN) 26 | poller.register(in_stream_sock, zmq.POLLIN) 27 | 28 | class Connection(object): 29 | def __init__(self, rid): 30 | self.rid = rid 31 | self.seq = 0 32 | self.exp_time = None 33 | 34 | def send_msg(self, msg): 35 | msg[b'from'] = instance_id 36 | msg[b'id'] = self.rid[1] 37 | msg[b'seq'] = self.seq 38 | self.seq += 1 39 | 40 | print('OUT {} {}'.format(self.rid[0], msg)) 41 | out_sock.send(self.rid[0] + b' T' + tnetstring.dumps(msg)) 42 | 43 | def send_header(self): 44 | msg = {} 45 | msg[b'code'] = 200 46 | msg[b'reason'] = b'OK' 47 | msg[b'headers'] = [[b'Content-Type', b'text/plain']] 48 | msg[b'more'] = True 49 | 50 | self.send_msg(msg) 51 | 52 | def send_body(self, data): 53 | msg = {} 54 | msg[b'body'] = data 55 | msg[b'more'] = True 56 | 57 | self.send_msg(msg) 58 | 59 | def send_body(to_addr, conns, data): 60 | ids = [] 61 | for c in conns: 62 | ids.append({b'id': c.rid[1], b'seq': c.seq}) 63 | c.seq += 1 64 | 65 | msg = {} 66 | msg[b'from'] = instance_id 67 | msg[b'id'] = ids 68 | msg[b'body'] = data 69 | msg[b'more'] = True 70 | 71 | print('OUT {} {}'.format(to_addr, msg)) 72 | out_sock.send(to_addr + b' T' + tnetstring.dumps(msg)) 73 | 74 | conns = {} 75 | last_exp_time = int(time.time()) 76 | 77 | while True: 78 | socks = dict(poller.poll(1000)) 79 | 80 | if socks.get(in_sock) == zmq.POLLIN: 81 | m_raw = in_sock.recv() 82 | elif socks.get(in_stream_sock) == zmq.POLLIN: 83 | m_list = in_stream_sock.recv_multipart() 84 | m_raw = m_list[2] 85 | else: 86 | m_raw = None 87 | 88 | now = int(time.time() * 1000) 89 | 90 | if m_raw is not None: 91 | req = tnetstring.loads(m_raw[1:]) 92 | print('IN {}'.format(req)) 93 | 94 | m_from = req[b'from'] 95 | m_id = req[b'id'] 96 | m_type = req.get(b'type', b'') 97 | 98 | ids = [] 99 | if isinstance(m_id, list): 100 | for id_seq in m_id: 101 | ids.append(id_seq[b'id']) 102 | else: 103 | ids.append(m_id) 104 | 105 | new_ids = [] 106 | known_conns = [] 107 | for i in ids: 108 | rid = (m_from, i) 109 | 110 | c = conns.get(rid) 111 | if c: 112 | c.exp_time = now + CONN_TTL 113 | known_conns.append(c) 114 | else: 115 | new_ids.append(rid) 116 | 117 | # data 118 | if not m_type: 119 | for rid in new_ids: 120 | c = Connection(rid) 121 | conns[rid] = c 122 | c.exp_time = now + CONN_TTL 123 | c.send_header() 124 | elif c: 125 | if m_type == b'keep-alive': 126 | dt = datetime.datetime.utcnow() 127 | ts = calendar.timegm(dt.timetuple()) 128 | 129 | body = ( 130 | 'id: TCPKaliMsgTS-{:016x}.\n' 131 | 'event: message\n' 132 | 'data: {:04}-{:02}-{:02}T{:02}:{:02}:{:02}\n\n' 133 | ).format( 134 | (ts * 1000000) + dt.microsecond, 135 | dt.year, 136 | dt.month, 137 | dt.day, 138 | dt.hour, 139 | dt.minute, 140 | dt.second 141 | ).encode() 142 | 143 | send_body(m_from, known_conns, body) 144 | elif m_type == b'cancel': 145 | for c in known_conns: 146 | del conns[c.rid] 147 | 148 | if now >= last_exp_time + EXPIRE_INTERVAL: 149 | last_exp_time = now 150 | 151 | to_remove = [] 152 | for rid, c in conns.items(): 153 | if now >= c.exp_time: 154 | to_remove.append(rid) 155 | for rid in to_remove: 156 | print('expired {}'.format(rid)) 157 | del conns[rid] 158 | -------------------------------------------------------------------------------- /benches/server.rs: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2020 Fanout, Inc. 3 | * 4 | * Licensed under the Apache License, Version 2.0 (the "License"); 5 | * you may not use this file except in compliance with the License. 6 | * You may obtain a copy of the License at 7 | * 8 | * http://www.apache.org/licenses/LICENSE-2.0 9 | * 10 | * Unless required by applicable law or agreed to in writing, software 11 | * distributed under the License is distributed on an "AS IS" BASIS, 12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | * See the License for the specific language governing permissions and 14 | * limitations under the License. 15 | */ 16 | 17 | use condure::connection::testutil::{ 18 | BenchServerReqConnection, BenchServerReqHandler, BenchServerStreamConnection, 19 | BenchServerStreamHandler, 20 | }; 21 | use condure::executor::Executor; 22 | use condure::future::{AsyncReadExt, AsyncTcpStream, AsyncWriteExt}; 23 | use condure::reactor::Reactor; 24 | use condure::server::TestServer; 25 | use condure::websocket::testutil::{BenchRecvMessage, BenchSendMessage}; 26 | use criterion::{criterion_group, criterion_main, Criterion}; 27 | use std::io::{self, Write}; 28 | use std::net::SocketAddr; 29 | use std::str; 30 | 31 | const REQS_PER_ITER: usize = 10; 32 | 33 | fn req(addr: SocketAddr) { 34 | let reactor = Reactor::new(REQS_PER_ITER * 10); 35 | let executor = Executor::new(REQS_PER_ITER); 36 | 37 | for _ in 0..REQS_PER_ITER { 38 | executor 39 | .spawn(async move { 40 | let mut client = AsyncTcpStream::connect(&[addr]).await.unwrap(); 41 | 42 | client 43 | .write(b"GET /hello HTTP/1.0\r\nHost: example.com\r\n\r\n") 44 | .await 45 | .unwrap(); 46 | 47 | let mut resp = [0u8; 1024]; 48 | let mut resp = io::Cursor::new(&mut resp[..]); 49 | 50 | loop { 51 | let mut buf = [0; 1024]; 52 | 53 | let size = client.read(&mut buf).await.unwrap(); 54 | if size == 0 { 55 | break; 56 | } 57 | 58 | resp.write(&buf[..size]).unwrap(); 59 | } 60 | 61 | let size = resp.position() as usize; 62 | let resp = str::from_utf8(&resp.get_ref()[..size]).unwrap(); 63 | 64 | assert_eq!(resp, "HTTP/1.0 200 OK\r\nContent-Length: 6\r\n\r\nworld\n"); 65 | }) 66 | .unwrap(); 67 | } 68 | 69 | executor.run(|timeout| reactor.poll(timeout)).unwrap(); 70 | } 71 | 72 | fn criterion_benchmark(c: &mut Criterion) { 73 | { 74 | let t = BenchServerReqHandler::new(); 75 | 76 | c.bench_function("req_handler", |b| { 77 | b.iter_batched_ref(|| t.init(), |i| t.run(i), criterion::BatchSize::SmallInput) 78 | }); 79 | } 80 | 81 | { 82 | let t = BenchServerStreamHandler::new(); 83 | 84 | c.bench_function("stream_handler", |b| { 85 | b.iter_batched_ref(|| t.init(), |i| t.run(i), criterion::BatchSize::SmallInput) 86 | }); 87 | } 88 | 89 | { 90 | let t = BenchServerReqConnection::new(); 91 | 92 | c.bench_function("req_connection", |b| { 93 | b.iter_batched_ref(|| t.init(), |i| t.run(i), criterion::BatchSize::SmallInput) 94 | }); 95 | } 96 | 97 | { 98 | let t = BenchServerStreamConnection::new(); 99 | 100 | c.bench_function("stream_connection", |b| { 101 | b.iter_batched_ref(|| t.init(), |i| t.run(i), criterion::BatchSize::SmallInput) 102 | }); 103 | } 104 | 105 | { 106 | let t = BenchSendMessage::new(false); 107 | 108 | c.bench_function("ws_send", |b| { 109 | b.iter_batched_ref(|| t.init(), |i| t.run(i), criterion::BatchSize::SmallInput) 110 | }); 111 | } 112 | 113 | { 114 | let t = BenchSendMessage::new(true); 115 | 116 | c.bench_function("ws_send_deflate", |b| { 117 | b.iter_batched_ref(|| t.init(), |i| t.run(i), criterion::BatchSize::SmallInput) 118 | }); 119 | } 120 | 121 | { 122 | let t = BenchRecvMessage::new(false); 123 | 124 | c.bench_function("ws_recv", |b| { 125 | b.iter_batched_ref(|| t.init(), |i| t.run(i), criterion::BatchSize::SmallInput) 126 | }); 127 | } 128 | 129 | { 130 | let t = BenchRecvMessage::new(true); 131 | 132 | c.bench_function("ws_recv_deflate", |b| { 133 | b.iter_batched_ref(|| t.init(), |i| t.run(i), criterion::BatchSize::SmallInput) 134 | }); 135 | } 136 | 137 | { 138 | let server = TestServer::new(1); 139 | let req_addr = server.req_addr(); 140 | let stream_addr = server.stream_addr(); 141 | 142 | c.bench_function("req_server workers=1", |b| b.iter(|| req(req_addr))); 143 | c.bench_function("stream_server workers=1", |b| b.iter(|| req(stream_addr))); 144 | } 145 | 146 | { 147 | let server = TestServer::new(2); 148 | let req_addr = server.req_addr(); 149 | let stream_addr = server.stream_addr(); 150 | 151 | c.bench_function("req_server workers=2", |b| b.iter(|| req(req_addr))); 152 | c.bench_function("stream_server workers=2", |b| b.iter(|| req(stream_addr))); 153 | } 154 | } 155 | 156 | criterion_group!(benches, criterion_benchmark); 157 | criterion_main!(benches); 158 | -------------------------------------------------------------------------------- /src/pool.rs: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2023 Fanout, Inc. 3 | * 4 | * Licensed under the Apache License, Version 2.0 (the "License"); 5 | * you may not use this file except in compliance with the License. 6 | * You may obtain a copy of the License at 7 | * 8 | * http://www.apache.org/licenses/LICENSE-2.0 9 | * 10 | * Unless required by applicable law or agreed to in writing, software 11 | * distributed under the License is distributed on an "AS IS" BASIS, 12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | * See the License for the specific language governing permissions and 14 | * limitations under the License. 15 | */ 16 | 17 | use crate::list; 18 | use crate::timer::TimerWheel; 19 | use slab::Slab; 20 | use std::borrow::Borrow; 21 | use std::collections::HashMap; 22 | use std::hash::Hash; 23 | use std::time::{Duration, Instant}; 24 | 25 | const TICK_DURATION_MS: u64 = 10; 26 | 27 | fn duration_to_ticks_round_down(d: Duration) -> u64 { 28 | (d.as_millis() / (TICK_DURATION_MS as u128)) as u64 29 | } 30 | 31 | struct PoolItem { 32 | key: K, 33 | value: V, 34 | timer_id: usize, 35 | } 36 | 37 | pub struct Pool { 38 | nodes: Slab>>, 39 | by_key: HashMap, 40 | wheel: TimerWheel, 41 | start: Instant, 42 | current_ticks: u64, 43 | } 44 | 45 | impl Pool 46 | where 47 | K: Clone + Eq + Hash + PartialEq, 48 | { 49 | pub fn new(capacity: usize) -> Self { 50 | Self { 51 | nodes: Slab::with_capacity(capacity), 52 | by_key: HashMap::with_capacity(capacity), 53 | wheel: TimerWheel::new(capacity), 54 | start: Instant::now(), 55 | current_ticks: 0, 56 | } 57 | } 58 | 59 | pub fn add(&mut self, key: K, value: V, expires: Instant) -> Result<(), V> { 60 | if self.nodes.len() == self.nodes.capacity() { 61 | return Err(value); 62 | } 63 | 64 | let expires = self.get_ticks(expires); 65 | 66 | let nkey = { 67 | let entry = self.nodes.vacant_entry(); 68 | let nkey = entry.key(); 69 | 70 | let timer_id = self.wheel.add(expires, nkey).unwrap(); 71 | 72 | entry.insert(list::Node::new(PoolItem { 73 | key: key.clone(), 74 | value, 75 | timer_id, 76 | })); 77 | 78 | nkey 79 | }; 80 | 81 | let l = self.by_key.entry(key).or_default(); 82 | 83 | l.push_back(&mut self.nodes, nkey); 84 | 85 | Ok(()) 86 | } 87 | 88 | pub fn take(&mut self, key: &Q) -> Option 89 | where 90 | K: Borrow, 91 | Q: Hash + Eq + ?Sized, 92 | { 93 | let l = self.by_key.get_mut(key)?; 94 | 95 | let nkey = l.pop_front(&mut self.nodes)?; 96 | 97 | if l.is_empty() { 98 | self.by_key.remove(key); 99 | } 100 | 101 | let pi = self.nodes.remove(nkey).value; 102 | 103 | self.wheel.remove(pi.timer_id); 104 | 105 | Some(pi.value) 106 | } 107 | 108 | pub fn expire(&mut self, now: Instant) -> Option<(K, V)> { 109 | let ticks = self.get_ticks(now); 110 | 111 | if ticks > self.current_ticks { 112 | self.wheel.update(ticks); 113 | self.current_ticks = ticks; 114 | } 115 | 116 | let nkey = match self.wheel.take_expired() { 117 | Some((_, nkey)) => nkey, 118 | None => return None, 119 | }; 120 | 121 | let pi = &self.nodes[nkey].value; 122 | 123 | let l = self.by_key.get_mut(&pi.key).unwrap(); 124 | l.remove(&mut self.nodes, nkey); 125 | 126 | if l.is_empty() { 127 | let pi = &self.nodes[nkey].value; 128 | self.by_key.remove(&pi.key); 129 | } 130 | 131 | let pi = self.nodes.remove(nkey).value; 132 | 133 | Some((pi.key, pi.value)) 134 | } 135 | 136 | fn get_ticks(&self, t: Instant) -> u64 { 137 | let d = if t > self.start { 138 | t - self.start 139 | } else { 140 | Duration::from_millis(0) 141 | }; 142 | 143 | duration_to_ticks_round_down(d) 144 | } 145 | } 146 | 147 | #[cfg(test)] 148 | mod tests { 149 | use super::*; 150 | 151 | #[test] 152 | fn pool_add_take() { 153 | let mut pool = Pool::new(3); 154 | 155 | let now = Instant::now(); 156 | pool.add(1, "a", now).unwrap(); 157 | pool.add(1, "b", now).unwrap(); 158 | pool.add(2, "c", now).unwrap(); 159 | assert_eq!(pool.add(2, "d", now).is_ok(), false); 160 | 161 | assert_eq!(pool.take(&1), Some("a")); 162 | assert_eq!(pool.take(&2), Some("c")); 163 | assert_eq!(pool.take(&1), Some("b")); 164 | assert_eq!(pool.take(&2), None); 165 | } 166 | 167 | #[test] 168 | fn pool_expire() { 169 | let mut pool = Pool::new(3); 170 | 171 | let now = Instant::now(); 172 | pool.add(1, "a", now + Duration::from_secs(1)).unwrap(); 173 | pool.add(1, "b", now + Duration::from_secs(2)).unwrap(); 174 | pool.add(2, "c", now + Duration::from_secs(3)).unwrap(); 175 | 176 | assert_eq!(pool.expire(now), None); 177 | assert_eq!(pool.expire(now + Duration::from_secs(1)), Some((1, "a"))); 178 | assert_eq!(pool.expire(now + Duration::from_secs(1)), None); 179 | assert_eq!(pool.expire(now + Duration::from_secs(5)), Some((1, "b"))); 180 | assert_eq!(pool.expire(now + Duration::from_secs(5)), Some((2, "c"))); 181 | assert_eq!(pool.expire(now + Duration::from_secs(5)), None); 182 | } 183 | } 184 | -------------------------------------------------------------------------------- /benches/client.rs: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2023 Fanout, Inc. 3 | * 4 | * Licensed under the Apache License, Version 2.0 (the "License"); 5 | * you may not use this file except in compliance with the License. 6 | * You may obtain a copy of the License at 7 | * 8 | * http://www.apache.org/licenses/LICENSE-2.0 9 | * 10 | * Unless required by applicable law or agreed to in writing, software 11 | * distributed under the License is distributed on an "AS IS" BASIS, 12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | * See the License for the specific language governing permissions and 14 | * limitations under the License. 15 | */ 16 | 17 | use condure::channel; 18 | use condure::client::TestClient; 19 | use condure::executor::Executor; 20 | use condure::future::{AsyncReadExt, AsyncSender, AsyncTcpListener, AsyncTcpStream, AsyncWriteExt}; 21 | use condure::reactor::Reactor; 22 | use criterion::{criterion_group, criterion_main, Criterion}; 23 | use mio::net::TcpListener; 24 | use std::net::SocketAddr; 25 | use std::rc::Rc; 26 | use std::str; 27 | 28 | const REQS_PER_ITER: usize = 10; 29 | 30 | fn req(listener: TcpListener, start: F1, wait: F2) -> TcpListener 31 | where 32 | F1: Fn(SocketAddr) + 'static, 33 | F2: Fn() + 'static, 34 | { 35 | let executor = Executor::new(REQS_PER_ITER + 1); 36 | 37 | let addr = listener.local_addr().unwrap(); 38 | 39 | let (s, r) = channel::channel(1); 40 | 41 | for _ in 0..REQS_PER_ITER { 42 | start(addr); 43 | } 44 | 45 | let spawner = executor.spawner(); 46 | 47 | executor 48 | .spawn(async move { 49 | let s = AsyncSender::new(s); 50 | let listener = AsyncTcpListener::new(listener); 51 | 52 | for _ in 0..REQS_PER_ITER { 53 | let (stream, _) = listener.accept().await.unwrap(); 54 | let mut stream = AsyncTcpStream::new(stream); 55 | 56 | spawner 57 | .spawn(async move { 58 | let mut buf = Vec::new(); 59 | let mut req_end = 0; 60 | 61 | while req_end == 0 { 62 | let mut chunk = [0; 1024]; 63 | let size = stream.read(&mut chunk).await.unwrap(); 64 | buf.extend_from_slice(&chunk[..size]); 65 | 66 | for i in 0..(buf.len() - 3) { 67 | if &buf[i..(i + 4)] == b"\r\n\r\n" { 68 | req_end = i + 4; 69 | break; 70 | } 71 | } 72 | } 73 | 74 | let expected = format!( 75 | concat!("GET /path HTTP/1.1\r\n", "Host: {}\r\n", "\r\n"), 76 | addr 77 | ); 78 | 79 | assert_eq!(str::from_utf8(&buf[..req_end]).unwrap(), expected); 80 | 81 | stream 82 | .write( 83 | b"HTTP/1.1 200 OK\r\nConnection: close\r\nContent-Type: text/plain\r\nContent-Length: 6\r\n\r\nhello\n", 84 | ).await 85 | .unwrap(); 86 | }) 87 | .unwrap(); 88 | } 89 | 90 | s.send(listener.into_inner()).await.unwrap(); 91 | }) 92 | .unwrap(); 93 | 94 | executor 95 | .run(|timeout| Reactor::current().unwrap().poll(timeout)) 96 | .unwrap(); 97 | 98 | for _ in 0..REQS_PER_ITER { 99 | wait(); 100 | } 101 | 102 | let listener = r.recv().unwrap(); 103 | 104 | listener 105 | } 106 | 107 | fn criterion_benchmark(c: &mut Criterion) { 108 | let mut req_listener = Some(TcpListener::bind("127.0.0.1:0".parse().unwrap()).unwrap()); 109 | let mut stream_listener = Some(TcpListener::bind("127.0.0.1:0".parse().unwrap()).unwrap()); 110 | let _reactor = Reactor::new(REQS_PER_ITER * 10); 111 | 112 | { 113 | let client = Rc::new(TestClient::new(1)); 114 | 115 | c.bench_function("req_client workers=1", |b| { 116 | b.iter(|| { 117 | let c1 = Rc::clone(&client); 118 | let c2 = Rc::clone(&client); 119 | 120 | req_listener = Some(req( 121 | req_listener.take().unwrap(), 122 | move |addr| c1.do_req(addr), 123 | move || c2.wait_req(), 124 | )) 125 | }) 126 | }); 127 | c.bench_function("stream_client workers=1", |b| { 128 | b.iter(|| { 129 | let c1 = Rc::clone(&client); 130 | let c2 = Rc::clone(&client); 131 | 132 | stream_listener = Some(req( 133 | stream_listener.take().unwrap(), 134 | move |addr| c1.do_stream_http(addr), 135 | move || c2.wait_stream(), 136 | )) 137 | }) 138 | }); 139 | } 140 | 141 | { 142 | let client = Rc::new(TestClient::new(2)); 143 | 144 | c.bench_function("req_client workers=2", |b| { 145 | b.iter(|| { 146 | let c1 = Rc::clone(&client); 147 | let c2 = Rc::clone(&client); 148 | 149 | req_listener = Some(req( 150 | req_listener.take().unwrap(), 151 | move |addr| c1.do_req(addr), 152 | move || c2.wait_req(), 153 | )) 154 | }) 155 | }); 156 | c.bench_function("stream_client workers=2", |b| { 157 | b.iter(|| { 158 | let c1 = Rc::clone(&client); 159 | let c2 = Rc::clone(&client); 160 | 161 | stream_listener = Some(req( 162 | stream_listener.take().unwrap(), 163 | move |addr| c1.do_stream_http(addr), 164 | move || c2.wait_stream(), 165 | )) 166 | }) 167 | }); 168 | } 169 | } 170 | 171 | criterion_group!(benches, criterion_benchmark); 172 | criterion_main!(benches); 173 | -------------------------------------------------------------------------------- /README.md: -------------------------------------------------------------------------------- 1 | **UPDATE**: This project has been merged into [Pushpin](https://github.com/fastly/pushpin) and the repository has been archived. The `pushpin-connmgr` program can be used as a drop-in substitute. 2 | 3 | # Condure 4 | 5 | Condure is a service that manages network connections in order to allow controlling the connections from multiple processes. It can manage incoming connections as well as outgoing connections. Applications communicate with Condure over [ZeroMQ](https://zeromq.org/). 6 | 7 | Condure can only manage connections for protocols it knows about. Currently this is HTTP/1 and WebSockets. See [Supported protocols](#supported-protocols). 8 | 9 | The project was inspired by [Mongrel2](https://mongrel2.org/). 10 | 11 | ## Use cases 12 | 13 | * Pass connection ownership from one process to another. 14 | * Restart an application without its connections getting disconnected. 15 | * Balance connection ownership among multiple processes. 16 | 17 | ## Basic usage 18 | 19 | Start the server: 20 | 21 | ``` 22 | $ condure --listen 8000 --zclient-stream ipc://client 23 | ``` 24 | 25 | Connect a handler to it, such as this simple Python program: 26 | 27 | ```py 28 | # this handler responds to every request with "hello world" 29 | 30 | import os 31 | import time 32 | import tnetstring 33 | import zmq 34 | 35 | instance_id = 'basichandler.{}'.format(os.getpid()).encode() 36 | 37 | ctx = zmq.Context() 38 | in_sock = ctx.socket(zmq.PULL) 39 | in_sock.connect('ipc://client-out') 40 | out_sock = ctx.socket(zmq.PUB) 41 | out_sock.connect('ipc://client-in') 42 | 43 | # await subscription 44 | time.sleep(0.01) 45 | 46 | while True: 47 | m_raw = in_sock.recv() 48 | req = tnetstring.loads(m_raw[1:]) 49 | print('IN {}'.format(req)) 50 | 51 | resp = {} 52 | resp[b'from'] = instance_id 53 | resp[b'id'] = req[b'id'] 54 | resp[b'code'] = 200 55 | resp[b'reason'] = b'OK' 56 | resp[b'headers'] = [[b'Content-Type', b'text/plain']] 57 | resp[b'body'] = b'hello world\n' 58 | 59 | print('OUT {}'.format(resp)) 60 | out_sock.send(req[b'from'] + b' T' + tnetstring.dumps(resp)) 61 | ``` 62 | 63 | A client request: 64 | 65 | ``` 66 | $ curl -i http://localhost:8000 67 | HTTP/1.1 200 OK 68 | Content-Type: text/plain 69 | Content-Length: 12 70 | 71 | hello world 72 | ``` 73 | 74 | The process that receives the request doesn't need to be the same one that responds! For example, here's a program that outputs request IDs to stdout: 75 | 76 | ```py 77 | # this handler just outputs the request ID 78 | 79 | import tnetstring 80 | import zmq 81 | 82 | ctx = zmq.Context() 83 | sock = ctx.socket(zmq.PULL) 84 | sock.connect('ipc://client-out') 85 | 86 | while True: 87 | m = sock.recv_multipart() 88 | req = tnetstring.loads(m[0][1:]) 89 | print('{} {}'.format(req[b'from'].decode(), req[b'id'].decode())) 90 | ``` 91 | 92 | We can see request ID information when a client request is made: 93 | 94 | ``` 95 | $ python examples/printreq.py 96 | condure 0-0-0 97 | ``` 98 | 99 | From another shell we can respond using a program like this: 100 | 101 | ```py 102 | # this program sends a response to a certain request ID 103 | 104 | import sys 105 | import time 106 | import tnetstring 107 | import zmq 108 | 109 | body = sys.argv[1] 110 | addr = sys.argv[2].encode() 111 | rid = sys.argv[3].encode() 112 | 113 | ctx = zmq.Context() 114 | sock = ctx.socket(zmq.PUB) 115 | sock.connect('ipc://client-in') 116 | 117 | # await subscription 118 | time.sleep(0.01) 119 | 120 | resp = {} 121 | resp[b'from'] = b'sendresp' 122 | resp[b'id'] = rid 123 | resp[b'code'] = 200 124 | resp[b'reason'] = b'OK' 125 | resp[b'headers'] = [[b'Content-Type', b'text/plain']] 126 | resp[b'body'] = '{}\n'.format(body).encode() 127 | 128 | m = [addr + b' T' + tnetstring.dumps(resp)] 129 | 130 | sock.send_multipart(m) 131 | ``` 132 | 133 | For example: 134 | 135 | ``` 136 | $ python examples/sendresp.py "responding from another process" condure 0-0-0 137 | ``` 138 | 139 | The client sees: 140 | 141 | ``` 142 | $ curl -i http://localhost:8000 143 | HTTP/1.1 200 OK 144 | Content-Type: text/plain 145 | Content-Length: 32 146 | 147 | responding from another process 148 | ``` 149 | 150 | For easy testing, the programs can be piped together: 151 | 152 | ``` 153 | $ python -u examples/printreq.py | xargs -n 2 python examples/sendresp.py "responding from another process" 154 | ``` 155 | 156 | ## Suspending and resuming connections 157 | 158 | When passing control of a connection from one process to another, it is important to suspend the connection first. This is done by sending a `handoff-start` message and waiting for a `handoff-proceed` message. At that point, the connection information can be given to another process, and the connection can be resumed by sending any message (such as `keep-alive`). See the [ZHTTP spec](https://rfc.zeromq.org/spec/33/). 159 | 160 | ## REQ mode 161 | 162 | In addition to the stream mode which uses PUSH/ROUTER/SUB sockets, there is a "REQ" mode available which uses a DEALER socket. To enable it, set `req` as the mode on a listen port. This mode can be handy for implementing simple request/response servers using ZeroMQ. 163 | 164 | ## Supported protocols 165 | 166 | Condure supports HTTP/1 and WebSockets. 167 | 168 | Condure manages connections at layer 7 and only supports protocols it knows about. This is to simplify its usage. Handling arbitrary protocols would require applications to build protocol stacks capable of suspending/resuming sessions at arbitrary byte positions in TCP streams, making Condure usage prohibitive. Instead, Condure is protocol-aware, and provides parsed frames to applications, so that applications are only required to support suspending/resuming sessions at frame boundaries. 169 | 170 | ## Performance 171 | 172 | Condure was built for high performance. It uses numerous optimization techniques, including minimal heap allocations, ring buffers, vectored I/O, hierarchical timing wheels, and fast data structures (e.g. slabs). Over 1M concurrent connections have been tested on a single instance using just 2 workers (4 threads total). See https://blog.fanout.io/2020/08/11/rewriting-pushpins-connection-manager-in-rust/ 173 | 174 | ## Comparison to Mongrel2 175 | 176 | * Condure supports acting as a server and as a client. 177 | * Condure supports multiple cores. 178 | * Condure supports listening on multiple ports without requiring multiple processes. 179 | * Condure does not support multiple routes and is not intended to be a shared server. Each application that wants to keep connections in a separate process should spawn its own Condure instance. 180 | * Condure has no config file. Configuration is supplied using command line arguments. 181 | * Condure uses a different ZeroMQ-based protocol, [ZHTTP](https://rfc.zeromq.org/spec/33/), which is easier to use than Mongrel2's protocol and more reliable. 182 | 183 | ## Future plans 184 | 185 | * HTTP/2 186 | * HTTP/3 187 | -------------------------------------------------------------------------------- /src/lib.rs: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2020-2023 Fanout, Inc. 3 | * 4 | * Licensed under the Apache License, Version 2.0 (the "License"); 5 | * you may not use this file except in compliance with the License. 6 | * You may obtain a copy of the License at 7 | * 8 | * http://www.apache.org/licenses/LICENSE-2.0 9 | * 10 | * Unless required by applicable law or agreed to in writing, software 11 | * distributed under the License is distributed on an "AS IS" BASIS, 12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | * See the License for the specific language governing permissions and 14 | * limitations under the License. 15 | */ 16 | 17 | pub mod app; 18 | pub mod arena; 19 | pub mod buffer; 20 | pub mod channel; 21 | pub mod client; 22 | pub mod connection; 23 | pub mod event; 24 | pub mod executor; 25 | pub mod future; 26 | pub mod http1; 27 | pub mod list; 28 | pub mod listener; 29 | pub mod net; 30 | pub mod pool; 31 | pub mod reactor; 32 | pub mod resolver; 33 | pub mod server; 34 | pub mod shuffle; 35 | pub mod timer; 36 | pub mod tls; 37 | pub mod tnetstring; 38 | pub mod track; 39 | pub mod waker; 40 | pub mod websocket; 41 | pub mod zhttppacket; 42 | pub mod zhttpsocket; 43 | pub mod zmq; 44 | 45 | use app::Config; 46 | use log::info; 47 | use std::error::Error; 48 | use std::ffi::CString; 49 | use std::future::Future; 50 | use std::io; 51 | use std::mem; 52 | use std::ops::Deref; 53 | use std::os::unix::ffi::OsStrExt; 54 | use std::path::Path; 55 | use std::pin::Pin; 56 | use std::ptr; 57 | use std::task::{Context, Poll}; 58 | 59 | pub struct Defer { 60 | f: Option, 61 | } 62 | 63 | impl Defer { 64 | pub fn new(f: T) -> Self { 65 | Self { f: Some(f) } 66 | } 67 | } 68 | 69 | impl Drop for Defer { 70 | fn drop(&mut self) { 71 | let f = self.f.take().unwrap(); 72 | 73 | f(); 74 | } 75 | } 76 | 77 | pub struct Pinner<'a, T> { 78 | pub unsafe_pointer: &'a mut T, 79 | } 80 | 81 | impl<'a, T> Pinner<'a, T> { 82 | pub fn as_mut(&mut self) -> Pin<&mut T> { 83 | // SAFETY: as long as Pinner is only ever constructed via the pin!() 84 | // macro and the unsafe_pointer field is never directly accessed, 85 | // then the value is safe to pin here. this is because the macro 86 | // ensures the input is turned into a borrowed anonymous temporary, 87 | // preventing any further access to the original value after Pinner 88 | // is constructed, and Pinner has no methods that enable moving out 89 | // of the reference. the word "unsafe" is used in the field name to 90 | // discourage direct access. this is the best we can do, since the 91 | // field must be public for the macro to work. 92 | unsafe { Pin::new_unchecked(self.unsafe_pointer) } 93 | } 94 | 95 | pub fn set(&mut self, value: T) { 96 | self.as_mut().set(value) 97 | } 98 | } 99 | 100 | impl<'a, T> Pinner<'a, Option> { 101 | pub fn as_pin_mut(&mut self) -> Option> { 102 | self.as_mut().as_pin_mut() 103 | } 104 | } 105 | 106 | impl<'a, T> Deref for Pinner<'a, T> { 107 | type Target = T; 108 | 109 | fn deref(&self) -> &Self::Target { 110 | self.unsafe_pointer 111 | } 112 | } 113 | 114 | impl<'a, T> Future for Pinner<'a, T> 115 | where 116 | T: Future, 117 | { 118 | type Output = T::Output; 119 | 120 | fn poll(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { 121 | T::poll(Pin::into_inner(self).as_mut(), cx) 122 | } 123 | } 124 | 125 | // NOTE: replace with std::pin::pin someday 126 | #[macro_export] 127 | macro_rules! pin { 128 | ($x:expr) => { 129 | $crate::Pinner { 130 | unsafe_pointer: &mut { $x }, 131 | } 132 | }; 133 | } 134 | 135 | fn try_with_increasing_buffer(starting_size: usize, f: T) -> Result 136 | where 137 | T: Fn(&mut [u8]) -> Result, 138 | { 139 | let mut buf = Vec::new(); 140 | buf.resize(starting_size, 0); 141 | 142 | loop { 143 | match f(&mut buf) { 144 | Ok(v) => return Ok(v), 145 | Err(e) if e.raw_os_error() == Some(libc::ERANGE) => buf.resize(buf.len() * 2, 0), 146 | Err(e) => return Err(e), 147 | } 148 | } 149 | } 150 | 151 | fn get_user_uid(name: &str) -> Result { 152 | let name = CString::new(name).unwrap(); 153 | 154 | try_with_increasing_buffer(1024, |buf| unsafe { 155 | let mut pwd = mem::MaybeUninit::uninit(); 156 | let mut passwd = ptr::null_mut(); 157 | 158 | if libc::getpwnam_r( 159 | name.as_ptr(), 160 | pwd.as_mut_ptr(), 161 | buf.as_mut_ptr() as *mut i8, 162 | buf.len(), 163 | &mut passwd, 164 | ) != 0 165 | { 166 | return Err(io::Error::last_os_error()); 167 | } 168 | 169 | let passwd = match passwd.as_ref() { 170 | Some(r) => r, 171 | None => return Err(io::Error::from(io::ErrorKind::NotFound)), 172 | }; 173 | 174 | Ok(passwd.pw_uid) 175 | }) 176 | } 177 | 178 | fn get_group_gid(name: &str) -> Result { 179 | let name = CString::new(name).unwrap(); 180 | 181 | try_with_increasing_buffer(1024, |buf| unsafe { 182 | let mut grp = mem::MaybeUninit::uninit(); 183 | let mut group = ptr::null_mut(); 184 | 185 | if libc::getgrnam_r( 186 | name.as_ptr(), 187 | grp.as_mut_ptr(), 188 | buf.as_mut_ptr() as *mut i8, 189 | buf.len(), 190 | &mut group, 191 | ) != 0 192 | { 193 | return Err(io::Error::last_os_error()); 194 | } 195 | 196 | let group = match group.as_ref() { 197 | Some(r) => r, 198 | None => return Err(io::Error::from(io::ErrorKind::NotFound)), 199 | }; 200 | 201 | Ok(group.gr_gid) 202 | }) 203 | } 204 | 205 | pub fn set_user(path: &Path, user: &str) -> Result<(), io::Error> { 206 | let uid = get_user_uid(user)?; 207 | 208 | unsafe { 209 | let path = CString::new(path.as_os_str().as_bytes()).unwrap(); 210 | 211 | if libc::chown(path.as_ptr(), uid, u32::MAX) != 0 { 212 | return Err(io::Error::last_os_error()); 213 | } 214 | } 215 | 216 | Ok(()) 217 | } 218 | 219 | pub fn set_group(path: &Path, group: &str) -> Result<(), io::Error> { 220 | let gid = get_group_gid(group)?; 221 | 222 | unsafe { 223 | let path = CString::new(path.as_os_str().as_bytes()).unwrap(); 224 | 225 | if libc::chown(path.as_ptr(), u32::MAX, gid) != 0 { 226 | return Err(io::Error::last_os_error()); 227 | } 228 | } 229 | 230 | Ok(()) 231 | } 232 | 233 | pub fn can_move_mio_sockets_between_threads() -> bool { 234 | // on unix platforms, mio always uses epoll or kqueue, which support 235 | // this. mio makes no guarantee about supporting this on non-unix 236 | // platforms 237 | cfg!(unix) 238 | } 239 | 240 | pub fn run(config: &Config) -> Result<(), Box> { 241 | info!("starting..."); 242 | 243 | { 244 | let a = match app::App::new(config) { 245 | Ok(a) => a, 246 | Err(e) => { 247 | return Err(e.into()); 248 | } 249 | }; 250 | 251 | info!("started"); 252 | 253 | a.wait_for_term(); 254 | 255 | info!("stopping..."); 256 | } 257 | 258 | info!("stopped"); 259 | 260 | Ok(()) 261 | } 262 | -------------------------------------------------------------------------------- /src/track.rs: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2023 Fanout, Inc. 3 | * 4 | * Licensed under the Apache License, Version 2.0 (the "License"); 5 | * you may not use this file except in compliance with the License. 6 | * You may obtain a copy of the License at 7 | * 8 | * http://www.apache.org/licenses/LICENSE-2.0 9 | * 10 | * Unless required by applicable law or agreed to in writing, software 11 | * distributed under the License is distributed on an "AS IS" BASIS, 12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | * See the License for the specific language governing permissions and 14 | * limitations under the License. 15 | */ 16 | 17 | use crate::future::AsyncLocalReceiver; 18 | use std::cell::Cell; 19 | use std::future::Future; 20 | use std::ops::Deref; 21 | use std::pin::Pin; 22 | use std::sync::mpsc; 23 | use std::task::{Context, Poll}; 24 | 25 | #[derive(Default)] 26 | pub struct TrackFlag(Cell); 27 | 28 | impl TrackFlag { 29 | pub fn get(&self) -> bool { 30 | self.0.get() 31 | } 32 | 33 | pub fn set(&self, v: bool) { 34 | self.0.set(v); 35 | } 36 | } 37 | 38 | struct TrackInner<'a, T> { 39 | value: T, 40 | active: &'a TrackFlag, 41 | } 42 | 43 | // wrap a value and a shared flag representing the value's liveness. on init, 44 | // the flag is set to true. on drop, the flag is set to false 45 | pub struct Track<'a, T> { 46 | inner: Option>, 47 | } 48 | 49 | impl<'a, T> Track<'a, T> { 50 | pub fn new(value: T, active: &'a TrackFlag) -> Self { 51 | active.set(true); 52 | 53 | Self { 54 | inner: Some(TrackInner { value, active }), 55 | } 56 | } 57 | } 58 | 59 | impl<'a, A, B> Track<'a, (A, B)> { 60 | pub fn map_first(mut orig: Self) -> (Track<'a, A>, B) { 61 | let ((a, b), active) = { 62 | let inner = orig.inner.take().unwrap(); 63 | drop(orig); 64 | 65 | (inner.value, inner.active) 66 | }; 67 | 68 | (Track::new(a, active), b) 69 | } 70 | } 71 | 72 | impl<'a, T> Drop for Track<'a, T> { 73 | fn drop(&mut self) { 74 | if let Some(inner) = &self.inner { 75 | inner.active.set(false); 76 | } 77 | } 78 | } 79 | 80 | impl<'a, T> Deref for Track<'a, T> { 81 | type Target = T; 82 | 83 | fn deref(&self) -> &Self::Target { 84 | &self.inner.as_ref().unwrap().value 85 | } 86 | } 87 | 88 | // wrap an AsyncLocalReceiver and a shared flag representing the liveness of 89 | // one received value at a time. each received value is wrapped in Track and 90 | // must be dropped before reading the next value 91 | pub struct TrackedAsyncLocalReceiver<'a, T> { 92 | inner: AsyncLocalReceiver, 93 | value_active: &'a TrackFlag, 94 | } 95 | 96 | impl<'a, T> TrackedAsyncLocalReceiver<'a, T> { 97 | pub fn new(r: AsyncLocalReceiver, value_active: &'a TrackFlag) -> Self { 98 | value_active.set(false); 99 | 100 | Self { 101 | inner: r, 102 | value_active, 103 | } 104 | } 105 | 106 | // attempt to receive a value from the inner receiver. if a previously 107 | // received value has not been dropped, this method returns an error 108 | pub async fn recv(&self) -> Result, mpsc::RecvError> { 109 | if self.value_active.get() { 110 | return Err(mpsc::RecvError); 111 | } 112 | 113 | let v = self.inner.recv().await?; 114 | 115 | Ok(Track::new(v, self.value_active)) 116 | } 117 | } 118 | 119 | #[derive(Debug, PartialEq)] 120 | pub struct ValueActiveError; 121 | 122 | pub struct TrackFuture<'a, F> { 123 | fut: F, 124 | value_active: &'a TrackFlag, 125 | } 126 | 127 | impl<'a, F, T, E> Future for TrackFuture<'a, F> 128 | where 129 | F: Future>, 130 | E: From, 131 | { 132 | type Output = Result; 133 | 134 | fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll { 135 | // SAFETY: pin projection 136 | let (fut, value_active) = unsafe { 137 | let s = self.get_unchecked_mut(); 138 | let fut = Pin::new_unchecked(&mut s.fut); 139 | 140 | (fut, &s.value_active) 141 | }; 142 | 143 | let result = fut.poll(cx); 144 | 145 | if value_active.get() { 146 | return Poll::Ready(Err(ValueActiveError.into())); 147 | } 148 | 149 | result 150 | } 151 | } 152 | 153 | // wrap a future and a shared flag representing the liveness of some value. 154 | // if the value is true after polling the inner future, return an error 155 | pub fn track_future(fut: F, value_active: &TrackFlag) -> TrackFuture<'_, F> 156 | where 157 | F: Future, 158 | { 159 | TrackFuture { fut, value_active } 160 | } 161 | 162 | #[cfg(test)] 163 | mod tests { 164 | use super::*; 165 | use crate::channel; 166 | use crate::executor::Executor; 167 | use crate::future::yield_task; 168 | use crate::reactor::Reactor; 169 | 170 | #[test] 171 | fn track_value() { 172 | let f = TrackFlag::default(); 173 | 174 | let v = Track::new(42, &f); 175 | assert!(f.get()); 176 | assert_eq!(*v, 42); 177 | 178 | drop(v); 179 | assert!(!f.get()); 180 | } 181 | 182 | #[test] 183 | fn track_async_local_receiver() { 184 | let reactor = Reactor::new(2); 185 | let executor = Executor::new(1); 186 | 187 | let (s, r) = channel::local_channel(2, 1, &reactor.local_registration_memory()); 188 | 189 | s.try_send(1).unwrap(); 190 | s.try_send(2).unwrap(); 191 | drop(s); 192 | 193 | executor 194 | .spawn(async move { 195 | let f = TrackFlag::default(); 196 | 197 | let r = TrackedAsyncLocalReceiver::new(AsyncLocalReceiver::new(r), &f); 198 | 199 | let v = r.recv().await.unwrap(); 200 | assert_eq!(*v, 1); 201 | assert!(r.recv().await.is_err()); 202 | 203 | drop(v); 204 | let v = r.recv().await.unwrap(); 205 | assert_eq!(*v, 2); 206 | assert!(r.recv().await.is_err()); 207 | 208 | // no values left 209 | drop(v); 210 | assert!(r.recv().await.is_err()); 211 | }) 212 | .unwrap(); 213 | 214 | executor.run(|timeout| reactor.poll(timeout)).unwrap(); 215 | } 216 | 217 | #[test] 218 | fn track_value_and_future() { 219 | let executor = Executor::new(1); 220 | 221 | executor 222 | .spawn(async move { 223 | let f = TrackFlag::default(); 224 | 225 | // awaiting while the flag is active is an error 226 | let ret = track_future( 227 | async { 228 | let v = Track::new(1, &f); 229 | 230 | // this line will cause the error 231 | yield_task().await; 232 | 233 | // this line never reached 234 | drop(v); 235 | 236 | Ok(()) 237 | }, 238 | &f, 239 | ) 240 | .await; 241 | assert_eq!(ret, Err(ValueActiveError)); 242 | 243 | // awaiting while the flag is not active is ok 244 | let ret: Result<_, ValueActiveError> = track_future( 245 | async { 246 | let v = Track::new(1, &f); 247 | drop(v); 248 | 249 | // this is ok 250 | yield_task().await; 251 | 252 | Ok(()) 253 | }, 254 | &f, 255 | ) 256 | .await; 257 | assert_eq!(ret, Ok(())); 258 | }) 259 | .unwrap(); 260 | 261 | executor.run(|_| Ok(())).unwrap(); 262 | } 263 | } 264 | -------------------------------------------------------------------------------- /src/waker.rs: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2021 Fanout, Inc. 3 | * 4 | * Licensed under the Apache License, Version 2.0 (the "License"); 5 | * you may not use this file except in compliance with the License. 6 | * You may obtain a copy of the License at 7 | * 8 | * http://www.apache.org/licenses/LICENSE-2.0 9 | * 10 | * Unless required by applicable law or agreed to in writing, software 11 | * distributed under the License is distributed on an "AS IS" BASIS, 12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | * See the License for the specific language governing permissions and 14 | * limitations under the License. 15 | */ 16 | 17 | use std::cell::Cell; 18 | use std::mem; 19 | use std::ops::Deref; 20 | use std::rc::Rc; 21 | use std::task::{RawWaker, RawWakerVTable, Waker}; 22 | 23 | // adapted from std::task::Wake 24 | pub trait RcWake { 25 | fn wake(self: Rc); 26 | 27 | fn wake_by_ref(self: &Rc) { 28 | self.clone().wake(); 29 | } 30 | } 31 | 32 | pub fn into_std(waker: Rc) -> Waker { 33 | // SAFETY: This is safe because raw_waker safely constructs 34 | // a RawWaker from Rc. 35 | unsafe { Waker::from_raw(raw_waker(waker)) } 36 | } 37 | 38 | #[inline(always)] 39 | fn raw_waker(waker: Rc) -> RawWaker { 40 | unsafe fn clone_waker(waker: *const ()) -> RawWaker { 41 | let waker = mem::ManuallyDrop::new(Rc::from_raw(waker as *const W)); 42 | 43 | let waker = Rc::clone(&waker); 44 | 45 | RawWaker::new( 46 | Rc::into_raw(waker) as *const (), 47 | &RawWakerVTable::new( 48 | clone_waker::, 49 | wake::, 50 | wake_by_ref::, 51 | drop_waker::, 52 | ), 53 | ) 54 | } 55 | 56 | unsafe fn wake(waker: *const ()) { 57 | let waker = Rc::from_raw(waker as *const W); 58 | ::wake(waker); 59 | } 60 | 61 | unsafe fn wake_by_ref(waker: *const ()) { 62 | let waker = mem::ManuallyDrop::new(Rc::from_raw(waker as *const W)); 63 | ::wake_by_ref(&waker); 64 | } 65 | 66 | unsafe fn drop_waker(waker: *const ()) { 67 | Rc::from_raw(waker as *const W); 68 | } 69 | 70 | RawWaker::new( 71 | Rc::into_raw(waker) as *const (), 72 | &RawWakerVTable::new( 73 | clone_waker::, 74 | wake::, 75 | wake_by_ref::, 76 | drop_waker::, 77 | ), 78 | ) 79 | } 80 | 81 | pub trait RefWake { 82 | fn wake(&self); 83 | } 84 | 85 | pub struct RefWakerData { 86 | refs: Cell, 87 | w: W, 88 | } 89 | 90 | impl RefWakerData { 91 | pub fn new(w: W) -> Self { 92 | Self { 93 | refs: Cell::new(1), 94 | w, 95 | } 96 | } 97 | } 98 | 99 | // a waker that borrows its inner data and panics on drop if there are any 100 | // active clones. this makes it possible to create wakers without heap 101 | // allocations, as long as all clones are cleaned up before the RefWaker is 102 | // dropped 103 | pub struct RefWaker<'a, W: RefWake + 'a> { 104 | data: &'a RefWakerData, 105 | } 106 | 107 | impl<'a, W: RefWake + 'a> RefWaker<'a, W> { 108 | pub fn new(data: &'a RefWakerData) -> Self { 109 | Self { data } 110 | } 111 | 112 | pub fn ref_count(&self) -> usize { 113 | self.data.refs.get() 114 | } 115 | 116 | pub fn as_std<'b>(&self, scratch: &'b mut mem::MaybeUninit) -> &'b Waker { 117 | let rw = RawWaker::new( 118 | self.data as *const RefWakerData as *const (), 119 | &RawWakerVTable::new(Self::clone, Self::wake, Self::wake_by_ref, Self::drop), 120 | ); 121 | 122 | // SAFETY: the inner Waker data is part of self, and our Drop impl 123 | // guarantees any Waker instances cannot outlive self 124 | let waker = unsafe { Waker::from_raw(rw) }; 125 | 126 | scratch.write(waker); 127 | 128 | // SAFETY: scratch is initialized above 129 | unsafe { scratch.assume_init_ref() } 130 | } 131 | 132 | unsafe fn clone(data: *const ()) -> RawWaker { 133 | let data_ref = (data as *const RefWakerData).as_ref().unwrap(); 134 | 135 | data_ref.refs.set(data_ref.refs.get() + 1); 136 | 137 | RawWaker::new( 138 | data, 139 | &RawWakerVTable::new(Self::clone, Self::wake, Self::wake_by_ref, Self::drop), 140 | ) 141 | } 142 | 143 | unsafe fn wake(data: *const ()) { 144 | Self::wake_by_ref(data); 145 | 146 | Self::drop(data); 147 | } 148 | 149 | unsafe fn wake_by_ref(data: *const ()) { 150 | let data_ref = (data as *const RefWakerData).as_ref().unwrap(); 151 | 152 | data_ref.w.wake(); 153 | } 154 | 155 | unsafe fn drop(data: *const ()) { 156 | let data_ref = (data as *const RefWakerData).as_ref().unwrap(); 157 | 158 | let refs = data_ref.refs.get(); 159 | 160 | assert!(refs > 1); 161 | 162 | data_ref.refs.set(refs - 1); 163 | } 164 | } 165 | 166 | impl<'a, W: RefWake + 'a> Drop for RefWaker<'a, W> { 167 | fn drop(&mut self) { 168 | assert_eq!(self.data.refs.get(), 1); 169 | } 170 | } 171 | 172 | impl<'a, W: RefWake + 'a> Deref for RefWaker<'a, W> { 173 | type Target = W; 174 | 175 | fn deref(&self) -> &Self::Target { 176 | &self.data.w 177 | } 178 | } 179 | 180 | #[cfg(test)] 181 | mod tests { 182 | use super::*; 183 | use std::cell::Cell; 184 | 185 | struct TestWaker { 186 | waked: Cell, 187 | } 188 | 189 | impl TestWaker { 190 | fn new() -> Self { 191 | TestWaker { 192 | waked: Cell::new(0), 193 | } 194 | } 195 | 196 | fn waked(&self) -> u32 { 197 | self.waked.get() 198 | } 199 | } 200 | 201 | impl RcWake for TestWaker { 202 | fn wake(self: Rc) { 203 | self.waked.set(self.waked.get() + 1); 204 | } 205 | } 206 | 207 | impl RefWake for TestWaker { 208 | fn wake(&self) { 209 | self.waked.set(self.waked.get() + 1); 210 | } 211 | } 212 | 213 | #[test] 214 | fn test_rc_waker() { 215 | let data = Rc::new(TestWaker::new()); 216 | 217 | assert_eq!(Rc::strong_count(&data), 1); 218 | 219 | let waker = into_std(data.clone()); 220 | 221 | assert_eq!(Rc::strong_count(&data), 2); 222 | 223 | let waker2 = waker.clone(); 224 | 225 | assert_eq!(Rc::strong_count(&data), 3); 226 | assert_eq!(data.waked(), 0); 227 | 228 | waker2.wake(); 229 | 230 | assert_eq!(Rc::strong_count(&data), 2); 231 | assert_eq!(data.waked(), 1); 232 | 233 | waker.wake_by_ref(); 234 | 235 | assert_eq!(Rc::strong_count(&data), 2); 236 | assert_eq!(data.waked(), 2); 237 | 238 | mem::drop(waker); 239 | 240 | assert_eq!(Rc::strong_count(&data), 1); 241 | } 242 | 243 | #[test] 244 | fn test_ref_waker() { 245 | let data = RefWakerData::new(TestWaker::new()); 246 | let data = RefWaker::new(&data); 247 | 248 | assert_eq!(RefWaker::ref_count(&data), 1); 249 | 250 | let waker = data.as_std(&mut mem::MaybeUninit::uninit()).clone(); 251 | 252 | assert_eq!(RefWaker::ref_count(&data), 2); 253 | 254 | let waker2 = waker.clone(); 255 | 256 | assert_eq!(RefWaker::ref_count(&data), 3); 257 | assert_eq!(data.waked(), 0); 258 | 259 | waker2.wake(); 260 | 261 | assert_eq!(RefWaker::ref_count(&data), 2); 262 | assert_eq!(data.waked(), 1); 263 | 264 | waker.wake_by_ref(); 265 | 266 | assert_eq!(RefWaker::ref_count(&data), 2); 267 | assert_eq!(data.waked(), 2); 268 | 269 | mem::drop(waker); 270 | 271 | assert_eq!(RefWaker::ref_count(&data), 1); 272 | } 273 | } 274 | -------------------------------------------------------------------------------- /src/listener.rs: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2020-2022 Fanout, Inc. 3 | * 4 | * Licensed under the Apache License, Version 2.0 (the "License"); 5 | * you may not use this file except in compliance with the License. 6 | * You may obtain a copy of the License at 7 | * 8 | * http://www.apache.org/licenses/LICENSE-2.0 9 | * 10 | * Unless required by applicable law or agreed to in writing, software 11 | * distributed under the License is distributed on an "AS IS" BASIS, 12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | * See the License for the specific language governing permissions and 14 | * limitations under the License. 15 | */ 16 | 17 | use crate::arena::recycle_vec; 18 | use crate::channel; 19 | use crate::executor::Executor; 20 | use crate::future::{ 21 | select_2, select_slice, AsyncNetListener, AsyncReceiver, AsyncSender, NetAcceptFuture, Select2, 22 | WaitWritableFuture, 23 | }; 24 | use crate::net::{NetListener, NetStream, SocketAddr}; 25 | use crate::reactor::Reactor; 26 | use log::{debug, error}; 27 | use std::cmp; 28 | use std::sync::mpsc; 29 | use std::thread; 30 | 31 | const REACTOR_REGISTRATIONS_MAX: usize = 128; 32 | const EXECUTOR_TASKS_MAX: usize = 1; 33 | 34 | pub struct Listener { 35 | thread: Option>, 36 | stop: channel::Sender<()>, 37 | } 38 | 39 | impl Listener { 40 | pub fn new( 41 | name: &str, 42 | listeners: Vec, 43 | senders: Vec>, 44 | ) -> Listener { 45 | let (s, r) = channel::channel(1); 46 | 47 | let thread = thread::Builder::new() 48 | .name(name.to_string()) 49 | .spawn(move || { 50 | let reactor = Reactor::new(REACTOR_REGISTRATIONS_MAX); 51 | let executor = Executor::new(EXECUTOR_TASKS_MAX); 52 | 53 | executor.spawn(Self::run(r, listeners, senders)).unwrap(); 54 | 55 | executor.run(|timeout| reactor.poll(timeout)).unwrap(); 56 | }) 57 | .unwrap(); 58 | 59 | Self { 60 | thread: Some(thread), 61 | stop: s, 62 | } 63 | } 64 | 65 | async fn run( 66 | stop: channel::Receiver<()>, 67 | listeners: Vec, 68 | senders: Vec>, 69 | ) { 70 | let stop = AsyncReceiver::new(stop); 71 | 72 | let mut listeners: Vec = 73 | listeners.into_iter().map(AsyncNetListener::new).collect(); 74 | 75 | let mut senders: Vec> = 76 | senders.into_iter().map(AsyncSender::new).collect(); 77 | 78 | let mut listeners_pos = 0; 79 | let mut senders_pos = 0; 80 | 81 | let mut sender_tasks_mem: Vec> = 82 | Vec::with_capacity(senders.len()); 83 | 84 | let mut listener_tasks_mem: Vec = Vec::with_capacity(listeners.len()); 85 | 86 | let mut slice_scratch = Vec::with_capacity(cmp::max(senders.len(), listeners.len())); 87 | 88 | let mut stop_recv = stop.recv(); 89 | 90 | 'accept: loop { 91 | // wait for a sender to become writable 92 | 93 | let mut sender_tasks = recycle_vec(sender_tasks_mem); 94 | 95 | for s in senders.iter_mut() { 96 | sender_tasks.push(s.wait_writable()); 97 | } 98 | 99 | let result = select_2( 100 | &mut stop_recv, 101 | select_slice(&mut sender_tasks, &mut slice_scratch), 102 | ) 103 | .await; 104 | 105 | sender_tasks_mem = recycle_vec(sender_tasks); 106 | 107 | match result { 108 | Select2::R1(_) => break, 109 | Select2::R2(_) => {} 110 | } 111 | 112 | // accept a connection 113 | 114 | let mut listener_tasks = recycle_vec(listener_tasks_mem); 115 | 116 | let (b, a) = listeners.split_at_mut(listeners_pos); 117 | 118 | for l in a.iter_mut().chain(b.iter_mut()) { 119 | listener_tasks.push(l.accept()); 120 | } 121 | 122 | let (pos, stream, peer_addr) = loop { 123 | match select_2( 124 | &mut stop_recv, 125 | select_slice(&mut listener_tasks, &mut slice_scratch), 126 | ) 127 | .await 128 | { 129 | Select2::R1(_) => break 'accept, 130 | Select2::R2((pos, result)) => match result { 131 | Ok((stream, peer_addr)) => break (pos, stream, peer_addr), 132 | Err(e) => error!("accept error: {:?}", e), 133 | }, 134 | } 135 | }; 136 | 137 | listener_tasks_mem = recycle_vec(listener_tasks); 138 | 139 | let pos = (listeners_pos + pos) % listeners.len(); 140 | 141 | debug!("accepted connection from {}", peer_addr); 142 | 143 | listeners_pos = (pos + 1) % listeners.len(); 144 | 145 | // write connection to sender 146 | 147 | let mut pending_sock = Some((pos, stream, peer_addr)); 148 | 149 | for _ in 0..senders.len() { 150 | let sender = &mut senders[senders_pos]; 151 | 152 | if !sender.is_writable() { 153 | senders_pos = (senders_pos + 1) % senders.len(); 154 | continue; 155 | } 156 | 157 | let s = pending_sock.take().unwrap(); 158 | 159 | match sender.try_send(s) { 160 | Ok(()) => {} 161 | Err(mpsc::TrySendError::Full(s)) => pending_sock = Some(s), 162 | Err(mpsc::TrySendError::Disconnected(_)) => { 163 | // this could happen during shutdown 164 | debug!("receiver disconnected"); 165 | } 166 | } 167 | 168 | senders_pos = (senders_pos + 1) % senders.len(); 169 | 170 | if pending_sock.is_none() { 171 | break; 172 | } 173 | } 174 | } 175 | } 176 | } 177 | 178 | impl Drop for Listener { 179 | fn drop(&mut self) { 180 | // this should never fail. receiver won't disconnect unless 181 | // we tell it to 182 | self.stop.send(()).unwrap(); 183 | 184 | let thread = self.thread.take().unwrap(); 185 | thread.join().unwrap(); 186 | } 187 | } 188 | 189 | #[cfg(test)] 190 | mod tests { 191 | use super::*; 192 | use crate::event; 193 | use mio::net::TcpListener; 194 | use std::io::{Read, Write}; 195 | use std::mem; 196 | use std::sync::mpsc; 197 | 198 | #[test] 199 | fn test_accept() { 200 | let mut addrs = Vec::new(); 201 | let mut listeners = Vec::new(); 202 | let mut senders = Vec::new(); 203 | let mut receivers = Vec::new(); 204 | 205 | for _ in 0..2 { 206 | let addr = "127.0.0.1:0".parse().unwrap(); 207 | let l = TcpListener::bind(addr).unwrap(); 208 | addrs.push(l.local_addr().unwrap()); 209 | listeners.push(NetListener::Tcp(l)); 210 | 211 | let (sender, receiver) = channel::channel(0); 212 | senders.push(sender); 213 | receivers.push(receiver); 214 | } 215 | 216 | let _l = Listener::new("listener-test", listeners, senders); 217 | 218 | let mut poller = event::Poller::new(1024).unwrap(); 219 | 220 | let mut client = std::net::TcpStream::connect(&addrs[0]).unwrap(); 221 | 222 | poller 223 | .register_custom( 224 | receivers[0].get_read_registration(), 225 | mio::Token(1), 226 | mio::Interest::READABLE, 227 | ) 228 | .unwrap(); 229 | 230 | let result = receivers[0].try_recv(); 231 | assert_eq!(result.is_err(), true); 232 | assert_eq!(result.unwrap_err(), mpsc::TryRecvError::Empty); 233 | 234 | loop { 235 | poller.poll(None).unwrap(); 236 | 237 | let mut done = false; 238 | for event in poller.iter_events() { 239 | match event.token() { 240 | mio::Token(1) => { 241 | assert_eq!(event.is_readable(), true); 242 | done = true; 243 | break; 244 | } 245 | _ => unreachable!(), 246 | } 247 | } 248 | 249 | if done { 250 | break; 251 | } 252 | } 253 | 254 | let (lnum, peer_client, _) = receivers[0].try_recv().unwrap(); 255 | 256 | assert_eq!(lnum, 0); 257 | 258 | let mut peer_client = match peer_client { 259 | NetStream::Tcp(s) => s, 260 | _ => unreachable!(), 261 | }; 262 | 263 | peer_client.write(b"hello").unwrap(); 264 | mem::drop(peer_client); 265 | 266 | let mut buf = Vec::new(); 267 | client.read_to_end(&mut buf).unwrap(); 268 | assert_eq!(&buf, b"hello"); 269 | } 270 | } 271 | -------------------------------------------------------------------------------- /src/list.rs: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2020-2023 Fanout, Inc. 3 | * 4 | * Licensed under the Apache License, Version 2.0 (the "License"); 5 | * you may not use this file except in compliance with the License. 6 | * You may obtain a copy of the License at 7 | * 8 | * http://www.apache.org/licenses/LICENSE-2.0 9 | * 10 | * Unless required by applicable law or agreed to in writing, software 11 | * distributed under the License is distributed on an "AS IS" BASIS, 12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | * See the License for the specific language governing permissions and 14 | * limitations under the License. 15 | */ 16 | 17 | use std::ops::IndexMut; 18 | 19 | pub struct Node { 20 | pub prev: Option, 21 | pub next: Option, 22 | pub value: T, 23 | } 24 | 25 | impl Node { 26 | pub fn new(value: T) -> Self { 27 | Self { 28 | prev: None, 29 | next: None, 30 | value, 31 | } 32 | } 33 | } 34 | 35 | #[derive(Default, Clone, Copy)] 36 | pub struct List { 37 | pub head: Option, 38 | pub tail: Option, 39 | } 40 | 41 | impl List { 42 | pub fn is_empty(&self) -> bool { 43 | self.head.is_none() 44 | } 45 | 46 | pub fn insert(&mut self, nodes: &mut S, after: Option, key: usize) 47 | where 48 | S: IndexMut>, 49 | { 50 | let next = if let Some(pkey) = after { 51 | let pn = &mut nodes[pkey]; 52 | 53 | let next = pn.next; 54 | pn.next = Some(key); 55 | 56 | let n = &mut nodes[key]; 57 | n.prev = Some(pkey); 58 | 59 | next 60 | } else { 61 | let next = self.head; 62 | self.head = Some(key); 63 | 64 | let n = &mut nodes[key]; 65 | n.prev = None; 66 | 67 | next 68 | }; 69 | 70 | let n = &mut nodes[key]; 71 | n.next = next; 72 | 73 | if let Some(nkey) = next { 74 | let nn = &mut nodes[nkey]; 75 | 76 | nn.prev = Some(key); 77 | } else { 78 | self.tail = Some(key); 79 | } 80 | } 81 | 82 | pub fn remove(&mut self, nodes: &mut S, key: usize) 83 | where 84 | S: IndexMut>, 85 | { 86 | let n = &mut nodes[key]; 87 | 88 | let prev = n.prev.take(); 89 | let next = n.next.take(); 90 | 91 | if let Some(pkey) = prev { 92 | let pn = &mut nodes[pkey]; 93 | pn.next = next; 94 | } 95 | 96 | if let Some(nkey) = next { 97 | let nn = &mut nodes[nkey]; 98 | nn.prev = prev; 99 | } 100 | 101 | if let Some(hkey) = self.head { 102 | if hkey == key { 103 | self.head = next; 104 | } 105 | } 106 | 107 | if let Some(tkey) = self.tail { 108 | if tkey == key { 109 | self.tail = prev; 110 | } 111 | } 112 | } 113 | 114 | pub fn pop_front(&mut self, nodes: &mut S) -> Option 115 | where 116 | S: IndexMut>, 117 | { 118 | match self.head { 119 | Some(key) => { 120 | self.remove(nodes, key); 121 | 122 | Some(key) 123 | } 124 | None => None, 125 | } 126 | } 127 | 128 | pub fn push_back(&mut self, nodes: &mut S, key: usize) 129 | where 130 | S: IndexMut>, 131 | { 132 | self.insert(nodes, self.tail, key); 133 | } 134 | 135 | pub fn concat(&mut self, nodes: &mut S, other: &mut Self) 136 | where 137 | S: IndexMut>, 138 | { 139 | if other.is_empty() { 140 | // nothing to do 141 | return; 142 | } 143 | 144 | // other is non-empty so this is guaranteed to succeed 145 | let hkey = other.head.unwrap(); 146 | 147 | let next = nodes[hkey].next; 148 | 149 | // since we're inserting after the tail, this will set next=None 150 | self.insert(nodes, self.tail, hkey); 151 | 152 | // restore the node's next key 153 | nodes[hkey].next = next; 154 | 155 | self.tail = other.tail; 156 | 157 | other.head = None; 158 | other.tail = None; 159 | } 160 | 161 | pub fn iter<'a, T, S>(&self, nodes: &'a S) -> ListIterator<'a, S> 162 | where 163 | S: IndexMut>, 164 | { 165 | ListIterator { 166 | nodes, 167 | next: self.head, 168 | } 169 | } 170 | } 171 | 172 | pub struct ListIterator<'a, S> { 173 | nodes: &'a S, 174 | next: Option, 175 | } 176 | 177 | impl<'a, T, S> Iterator for ListIterator<'a, S> 178 | where 179 | T: 'a, 180 | S: IndexMut>, 181 | { 182 | type Item = (usize, &'a T); 183 | 184 | fn next(&mut self) -> Option { 185 | if let Some(nkey) = self.next.take() { 186 | let n = &self.nodes[nkey]; 187 | self.next = n.next; 188 | 189 | Some((nkey, &n.value)) 190 | } else { 191 | None 192 | } 193 | } 194 | } 195 | 196 | #[cfg(test)] 197 | mod tests { 198 | use super::*; 199 | use slab::Slab; 200 | 201 | #[test] 202 | fn test_list_push_pop() { 203 | let mut nodes = Slab::new(); 204 | let n1 = nodes.insert(Node::new("n1")); 205 | let n2 = nodes.insert(Node::new("n2")); 206 | let n3 = nodes.insert(Node::new("n3")); 207 | 208 | // prevent unused warning on data field 209 | assert_eq!(nodes[n1].value, "n1"); 210 | 211 | assert_eq!(nodes[n1].prev, None); 212 | assert_eq!(nodes[n1].next, None); 213 | assert_eq!(nodes[n2].prev, None); 214 | assert_eq!(nodes[n2].next, None); 215 | assert_eq!(nodes[n3].prev, None); 216 | assert_eq!(nodes[n3].next, None); 217 | 218 | let mut l = List::default(); 219 | assert_eq!(l.is_empty(), true); 220 | assert_eq!(l.head, None); 221 | assert_eq!(l.tail, None); 222 | assert_eq!(l.pop_front(&mut nodes), None); 223 | 224 | l.push_back(&mut nodes, n1); 225 | assert_eq!(l.is_empty(), false); 226 | assert_eq!(l.head, Some(n1)); 227 | assert_eq!(l.tail, Some(n1)); 228 | assert_eq!(nodes[n1].prev, None); 229 | assert_eq!(nodes[n1].next, None); 230 | 231 | l.push_back(&mut nodes, n2); 232 | assert_eq!(l.is_empty(), false); 233 | assert_eq!(l.head, Some(n1)); 234 | assert_eq!(l.tail, Some(n2)); 235 | assert_eq!(nodes[n1].prev, None); 236 | assert_eq!(nodes[n1].next, Some(n2)); 237 | assert_eq!(nodes[n2].prev, Some(n1)); 238 | assert_eq!(nodes[n2].next, None); 239 | 240 | l.push_back(&mut nodes, n3); 241 | assert_eq!(l.is_empty(), false); 242 | assert_eq!(l.head, Some(n1)); 243 | assert_eq!(l.tail, Some(n3)); 244 | assert_eq!(nodes[n1].prev, None); 245 | assert_eq!(nodes[n1].next, Some(n2)); 246 | assert_eq!(nodes[n2].prev, Some(n1)); 247 | assert_eq!(nodes[n2].next, Some(n3)); 248 | assert_eq!(nodes[n3].prev, Some(n2)); 249 | assert_eq!(nodes[n3].next, None); 250 | 251 | let key = l.pop_front(&mut nodes); 252 | assert_eq!(key, Some(n1)); 253 | assert_eq!(l.is_empty(), false); 254 | assert_eq!(l.head, Some(n2)); 255 | assert_eq!(l.tail, Some(n3)); 256 | assert_eq!(nodes[n2].prev, None); 257 | assert_eq!(nodes[n2].next, Some(n3)); 258 | assert_eq!(nodes[n3].prev, Some(n2)); 259 | assert_eq!(nodes[n3].next, None); 260 | 261 | let key = l.pop_front(&mut nodes); 262 | assert_eq!(key, Some(n2)); 263 | assert_eq!(l.is_empty(), false); 264 | assert_eq!(l.head, Some(n3)); 265 | assert_eq!(l.tail, Some(n3)); 266 | assert_eq!(nodes[n3].prev, None); 267 | assert_eq!(nodes[n3].next, None); 268 | 269 | let key = l.pop_front(&mut nodes); 270 | assert_eq!(key, Some(n3)); 271 | assert_eq!(l.is_empty(), true); 272 | assert_eq!(l.head, None); 273 | assert_eq!(l.tail, None); 274 | 275 | assert_eq!(l.pop_front(&mut nodes), None); 276 | } 277 | 278 | #[test] 279 | fn test_remove() { 280 | let mut nodes = Slab::new(); 281 | let n1 = nodes.insert(Node::new("n1")); 282 | 283 | assert_eq!(nodes[n1].prev, None); 284 | assert_eq!(nodes[n1].next, None); 285 | 286 | let mut l = List::default(); 287 | assert_eq!(l.is_empty(), true); 288 | assert_eq!(l.head, None); 289 | assert_eq!(l.tail, None); 290 | 291 | l.push_back(&mut nodes, n1); 292 | assert_eq!(l.is_empty(), false); 293 | assert_eq!(l.head, Some(n1)); 294 | assert_eq!(l.tail, Some(n1)); 295 | assert_eq!(nodes[n1].prev, None); 296 | assert_eq!(nodes[n1].next, None); 297 | 298 | l.remove(&mut nodes, n1); 299 | assert_eq!(l.is_empty(), true); 300 | assert_eq!(l.head, None); 301 | assert_eq!(l.tail, None); 302 | assert_eq!(nodes[n1].prev, None); 303 | assert_eq!(nodes[n1].next, None); 304 | 305 | // already removed 306 | l.remove(&mut nodes, n1); 307 | assert_eq!(l.is_empty(), true); 308 | assert_eq!(l.head, None); 309 | assert_eq!(l.tail, None); 310 | assert_eq!(nodes[n1].prev, None); 311 | assert_eq!(nodes[n1].next, None); 312 | } 313 | 314 | #[test] 315 | fn test_list_concat() { 316 | let mut nodes = Slab::new(); 317 | let n1 = nodes.insert(Node::new("n1")); 318 | let n2 = nodes.insert(Node::new("n2")); 319 | 320 | let mut a = List::default(); 321 | let mut b = List::default(); 322 | 323 | a.concat(&mut nodes, &mut b); 324 | assert_eq!(a.is_empty(), true); 325 | assert_eq!(a.head, None); 326 | assert_eq!(a.tail, None); 327 | assert_eq!(b.is_empty(), true); 328 | assert_eq!(b.head, None); 329 | assert_eq!(b.tail, None); 330 | 331 | a.push_back(&mut nodes, n1); 332 | b.push_back(&mut nodes, n2); 333 | 334 | a.concat(&mut nodes, &mut b); 335 | assert_eq!(a.is_empty(), false); 336 | assert_eq!(a.head, Some(n1)); 337 | assert_eq!(a.tail, Some(n2)); 338 | assert_eq!(b.is_empty(), true); 339 | assert_eq!(b.head, None); 340 | assert_eq!(b.tail, None); 341 | assert_eq!(nodes[n1].prev, None); 342 | assert_eq!(nodes[n1].next, Some(n2)); 343 | assert_eq!(nodes[n2].prev, Some(n1)); 344 | assert_eq!(nodes[n2].next, None); 345 | } 346 | 347 | #[test] 348 | fn test_list_iter() { 349 | let mut nodes = Slab::new(); 350 | let n1 = nodes.insert(Node::new("n1")); 351 | let n2 = nodes.insert(Node::new("n2")); 352 | let n3 = nodes.insert(Node::new("n3")); 353 | 354 | let mut l = List::default(); 355 | l.push_back(&mut nodes, n1); 356 | l.push_back(&mut nodes, n2); 357 | l.push_back(&mut nodes, n3); 358 | 359 | let mut it = l.iter(&nodes); 360 | assert_eq!(it.next(), Some((n1, &"n1"))); 361 | assert_eq!(it.next(), Some((n2, &"n2"))); 362 | assert_eq!(it.next(), Some((n3, &"n3"))); 363 | assert_eq!(it.next(), None); 364 | } 365 | } 366 | -------------------------------------------------------------------------------- /COPYING: -------------------------------------------------------------------------------- 1 | 2 | Apache License 3 | Version 2.0, January 2004 4 | http://www.apache.org/licenses/ 5 | 6 | TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION 7 | 8 | 1. Definitions. 9 | 10 | "License" shall mean the terms and conditions for use, reproduction, 11 | and distribution as defined by Sections 1 through 9 of this document. 12 | 13 | "Licensor" shall mean the copyright owner or entity authorized by 14 | the copyright owner that is granting the License. 15 | 16 | "Legal Entity" shall mean the union of the acting entity and all 17 | other entities that control, are controlled by, or are under common 18 | control with that entity. For the purposes of this definition, 19 | "control" means (i) the power, direct or indirect, to cause the 20 | direction or management of such entity, whether by contract or 21 | otherwise, or (ii) ownership of fifty percent (50%) or more of the 22 | outstanding shares, or (iii) beneficial ownership of such entity. 23 | 24 | "You" (or "Your") shall mean an individual or Legal Entity 25 | exercising permissions granted by this License. 26 | 27 | "Source" form shall mean the preferred form for making modifications, 28 | including but not limited to software source code, documentation 29 | source, and configuration files. 30 | 31 | "Object" form shall mean any form resulting from mechanical 32 | transformation or translation of a Source form, including but 33 | not limited to compiled object code, generated documentation, 34 | and conversions to other media types. 35 | 36 | "Work" shall mean the work of authorship, whether in Source or 37 | Object form, made available under the License, as indicated by a 38 | copyright notice that is included in or attached to the work 39 | (an example is provided in the Appendix below). 40 | 41 | "Derivative Works" shall mean any work, whether in Source or Object 42 | form, that is based on (or derived from) the Work and for which the 43 | editorial revisions, annotations, elaborations, or other modifications 44 | represent, as a whole, an original work of authorship. For the purposes 45 | of this License, Derivative Works shall not include works that remain 46 | separable from, or merely link (or bind by name) to the interfaces of, 47 | the Work and Derivative Works thereof. 48 | 49 | "Contribution" shall mean any work of authorship, including 50 | the original version of the Work and any modifications or additions 51 | to that Work or Derivative Works thereof, that is intentionally 52 | submitted to Licensor for inclusion in the Work by the copyright owner 53 | or by an individual or Legal Entity authorized to submit on behalf of 54 | the copyright owner. For the purposes of this definition, "submitted" 55 | means any form of electronic, verbal, or written communication sent 56 | to the Licensor or its representatives, including but not limited to 57 | communication on electronic mailing lists, source code control systems, 58 | and issue tracking systems that are managed by, or on behalf of, the 59 | Licensor for the purpose of discussing and improving the Work, but 60 | excluding communication that is conspicuously marked or otherwise 61 | designated in writing by the copyright owner as "Not a Contribution." 62 | 63 | "Contributor" shall mean Licensor and any individual or Legal Entity 64 | on behalf of whom a Contribution has been received by Licensor and 65 | subsequently incorporated within the Work. 66 | 67 | 2. Grant of Copyright License. Subject to the terms and conditions of 68 | this License, each Contributor hereby grants to You a perpetual, 69 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 70 | copyright license to reproduce, prepare Derivative Works of, 71 | publicly display, publicly perform, sublicense, and distribute the 72 | Work and such Derivative Works in Source or Object form. 73 | 74 | 3. Grant of Patent License. Subject to the terms and conditions of 75 | this License, each Contributor hereby grants to You a perpetual, 76 | worldwide, non-exclusive, no-charge, royalty-free, irrevocable 77 | (except as stated in this section) patent license to make, have made, 78 | use, offer to sell, sell, import, and otherwise transfer the Work, 79 | where such license applies only to those patent claims licensable 80 | by such Contributor that are necessarily infringed by their 81 | Contribution(s) alone or by combination of their Contribution(s) 82 | with the Work to which such Contribution(s) was submitted. If You 83 | institute patent litigation against any entity (including a 84 | cross-claim or counterclaim in a lawsuit) alleging that the Work 85 | or a Contribution incorporated within the Work constitutes direct 86 | or contributory patent infringement, then any patent licenses 87 | granted to You under this License for that Work shall terminate 88 | as of the date such litigation is filed. 89 | 90 | 4. Redistribution. You may reproduce and distribute copies of the 91 | Work or Derivative Works thereof in any medium, with or without 92 | modifications, and in Source or Object form, provided that You 93 | meet the following conditions: 94 | 95 | (a) You must give any other recipients of the Work or 96 | Derivative Works a copy of this License; and 97 | 98 | (b) You must cause any modified files to carry prominent notices 99 | stating that You changed the files; and 100 | 101 | (c) You must retain, in the Source form of any Derivative Works 102 | that You distribute, all copyright, patent, trademark, and 103 | attribution notices from the Source form of the Work, 104 | excluding those notices that do not pertain to any part of 105 | the Derivative Works; and 106 | 107 | (d) If the Work includes a "NOTICE" text file as part of its 108 | distribution, then any Derivative Works that You distribute must 109 | include a readable copy of the attribution notices contained 110 | within such NOTICE file, excluding those notices that do not 111 | pertain to any part of the Derivative Works, in at least one 112 | of the following places: within a NOTICE text file distributed 113 | as part of the Derivative Works; within the Source form or 114 | documentation, if provided along with the Derivative Works; or, 115 | within a display generated by the Derivative Works, if and 116 | wherever such third-party notices normally appear. The contents 117 | of the NOTICE file are for informational purposes only and 118 | do not modify the License. You may add Your own attribution 119 | notices within Derivative Works that You distribute, alongside 120 | or as an addendum to the NOTICE text from the Work, provided 121 | that such additional attribution notices cannot be construed 122 | as modifying the License. 123 | 124 | You may add Your own copyright statement to Your modifications and 125 | may provide additional or different license terms and conditions 126 | for use, reproduction, or distribution of Your modifications, or 127 | for any such Derivative Works as a whole, provided Your use, 128 | reproduction, and distribution of the Work otherwise complies with 129 | the conditions stated in this License. 130 | 131 | 5. Submission of Contributions. Unless You explicitly state otherwise, 132 | any Contribution intentionally submitted for inclusion in the Work 133 | by You to the Licensor shall be under the terms and conditions of 134 | this License, without any additional terms or conditions. 135 | Notwithstanding the above, nothing herein shall supersede or modify 136 | the terms of any separate license agreement you may have executed 137 | with Licensor regarding such Contributions. 138 | 139 | 6. Trademarks. This License does not grant permission to use the trade 140 | names, trademarks, service marks, or product names of the Licensor, 141 | except as required for reasonable and customary use in describing the 142 | origin of the Work and reproducing the content of the NOTICE file. 143 | 144 | 7. Disclaimer of Warranty. Unless required by applicable law or 145 | agreed to in writing, Licensor provides the Work (and each 146 | Contributor provides its Contributions) on an "AS IS" BASIS, 147 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or 148 | implied, including, without limitation, any warranties or conditions 149 | of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A 150 | PARTICULAR PURPOSE. You are solely responsible for determining the 151 | appropriateness of using or redistributing the Work and assume any 152 | risks associated with Your exercise of permissions under this License. 153 | 154 | 8. Limitation of Liability. In no event and under no legal theory, 155 | whether in tort (including negligence), contract, or otherwise, 156 | unless required by applicable law (such as deliberate and grossly 157 | negligent acts) or agreed to in writing, shall any Contributor be 158 | liable to You for damages, including any direct, indirect, special, 159 | incidental, or consequential damages of any character arising as a 160 | result of this License or out of the use or inability to use the 161 | Work (including but not limited to damages for loss of goodwill, 162 | work stoppage, computer failure or malfunction, or any and all 163 | other commercial damages or losses), even if such Contributor 164 | has been advised of the possibility of such damages. 165 | 166 | 9. Accepting Warranty or Additional Liability. While redistributing 167 | the Work or Derivative Works thereof, You may choose to offer, 168 | and charge a fee for, acceptance of support, warranty, indemnity, 169 | or other liability obligations and/or rights consistent with this 170 | License. However, in accepting such obligations, You may act only 171 | on Your own behalf and on Your sole responsibility, not on behalf 172 | of any other Contributor, and only if You agree to indemnify, 173 | defend, and hold each Contributor harmless for any liability 174 | incurred by, or claims asserted against, such Contributor by reason 175 | of your accepting any such warranty or additional liability. 176 | 177 | END OF TERMS AND CONDITIONS 178 | 179 | APPENDIX: How to apply the Apache License to your work. 180 | 181 | To apply the Apache License to your work, attach the following 182 | boilerplate notice, with the fields enclosed by brackets "[]" 183 | replaced with your own identifying information. (Don't include 184 | the brackets!) The text should be enclosed in the appropriate 185 | comment syntax for the file format. We also recommend that a 186 | file or class name and description of purpose be included on the 187 | same "printed page" as the copyright notice for easier 188 | identification within third-party archives. 189 | 190 | Copyright [yyyy] [name of copyright owner] 191 | 192 | Licensed under the Apache License, Version 2.0 (the "License"); 193 | you may not use this file except in compliance with the License. 194 | You may obtain a copy of the License at 195 | 196 | http://www.apache.org/licenses/LICENSE-2.0 197 | 198 | Unless required by applicable law or agreed to in writing, software 199 | distributed under the License is distributed on an "AS IS" BASIS, 200 | WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 201 | See the License for the specific language governing permissions and 202 | limitations under the License. 203 | -------------------------------------------------------------------------------- /src/resolver.rs: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2022 Fanout, Inc. 3 | * 4 | * Licensed under the Apache License, Version 2.0 (the "License"); 5 | * you may not use this file except in compliance with the License. 6 | * You may obtain a copy of the License at 7 | * 8 | * http://www.apache.org/licenses/LICENSE-2.0 9 | * 10 | * Unless required by applicable law or agreed to in writing, software 11 | * distributed under the License is distributed on an "AS IS" BASIS, 12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | * See the License for the specific language governing permissions and 14 | * limitations under the License. 15 | */ 16 | 17 | use crate::event; 18 | use crate::list; 19 | use arrayvec::{ArrayString, ArrayVec}; 20 | use mio::Interest; 21 | use slab::Slab; 22 | use std::collections::VecDeque; 23 | use std::io; 24 | use std::net::{IpAddr, ToSocketAddrs}; 25 | use std::sync::atomic::{AtomicBool, Ordering}; 26 | use std::sync::{Arc, Condvar, Mutex}; 27 | use std::thread; 28 | 29 | pub const REGISTRATIONS_PER_QUERY: usize = 1; 30 | 31 | pub const ADDRS_MAX: usize = 16; 32 | 33 | pub type Hostname = ArrayString<255>; 34 | pub type Addrs = ArrayVec; 35 | 36 | fn std_resolve(host: &str) -> Result { 37 | match (host, 0).to_socket_addrs() { 38 | Ok(addrs) => Ok(addrs.take(ADDRS_MAX).map(|addr| addr.ip()).collect()), 39 | Err(e) => Err(e), 40 | } 41 | } 42 | 43 | struct QueryItem { 44 | host: Hostname, 45 | result: Option>, 46 | set_readiness: event::SetReadiness, 47 | invalidated: Option>, 48 | } 49 | 50 | struct QueriesInner { 51 | stop: bool, 52 | nodes: Slab>, 53 | next: list::List, 54 | registrations: VecDeque<(event::Registration, event::SetReadiness)>, 55 | invalidated_count: u32, 56 | } 57 | 58 | #[derive(Clone)] 59 | struct Queries { 60 | inner: Arc<(Mutex, Condvar)>, 61 | } 62 | 63 | impl Queries { 64 | fn new(queries_max: usize) -> Self { 65 | let mut registrations = VecDeque::with_capacity(queries_max); 66 | 67 | for _ in 0..registrations.capacity() { 68 | registrations.push_back(event::Registration::new()); 69 | } 70 | 71 | let inner = QueriesInner { 72 | stop: false, 73 | nodes: Slab::with_capacity(queries_max), 74 | next: list::List::default(), 75 | registrations, 76 | invalidated_count: 0, 77 | }; 78 | 79 | Self { 80 | inner: Arc::new((Mutex::new(inner), Condvar::new())), 81 | } 82 | } 83 | 84 | fn set_stop_flag(&self) { 85 | let (lock, cvar) = &*self.inner; 86 | 87 | let mut queries = lock.lock().unwrap(); 88 | queries.stop = true; 89 | 90 | cvar.notify_all(); 91 | } 92 | 93 | fn add(&self, host: &str) -> Result<(usize, event::Registration), ()> { 94 | let (lock, cvar) = &*self.inner; 95 | 96 | let queries = &mut *lock.lock().unwrap(); 97 | 98 | if queries.nodes.len() == queries.nodes.capacity() { 99 | return Err(()); 100 | } 101 | 102 | let (reg, sr) = queries.registrations.pop_back().unwrap(); 103 | 104 | let nkey = match Hostname::from(host) { 105 | Ok(host) => { 106 | let nkey = queries.nodes.insert(list::Node::new(QueryItem { 107 | host, 108 | result: None, 109 | set_readiness: sr, 110 | invalidated: None, 111 | })); 112 | 113 | queries.next.push_back(&mut queries.nodes, nkey); 114 | 115 | cvar.notify_one(); 116 | 117 | nkey 118 | } 119 | Err(_) => { 120 | sr.set_readiness(Interest::READABLE).unwrap(); 121 | 122 | queries.nodes.insert(list::Node::new(QueryItem { 123 | host: Hostname::new(), 124 | result: Some(Err(io::Error::from(io::ErrorKind::InvalidInput))), 125 | set_readiness: sr, 126 | invalidated: None, 127 | })) 128 | } 129 | }; 130 | 131 | Ok((nkey, reg)) 132 | } 133 | 134 | // block until a query is available, or stopped 135 | fn get_next(&self, invalidated: &Arc) -> Option<(usize, Hostname)> { 136 | let (lock, cvar) = &*self.inner; 137 | 138 | let mut queries_guard = lock.lock().unwrap(); 139 | 140 | loop { 141 | let queries = &mut *queries_guard; 142 | 143 | if queries.stop { 144 | return None; 145 | } 146 | 147 | if let Some(nkey) = queries.next.pop_front(&mut queries.nodes) { 148 | let qi = &mut queries.nodes[nkey].value; 149 | 150 | invalidated.store(false, Ordering::Relaxed); 151 | qi.invalidated = Some(invalidated.clone()); 152 | 153 | return Some((nkey, qi.host)); 154 | } 155 | 156 | queries_guard = cvar.wait(queries_guard).unwrap(); 157 | } 158 | } 159 | 160 | fn set_result( 161 | &self, 162 | item_key: usize, 163 | result: Result, 164 | invalidated: &AtomicBool, 165 | ) { 166 | let mut queries = self.inner.0.lock().unwrap(); 167 | 168 | if !invalidated.load(Ordering::Relaxed) { 169 | let qi = &mut queries.nodes[item_key].value; 170 | 171 | qi.result = Some(result); 172 | qi.invalidated = None; 173 | qi.set_readiness.set_readiness(Interest::READABLE).unwrap(); 174 | } else { 175 | queries.invalidated_count += 1; 176 | } 177 | } 178 | 179 | fn take_result(&self, item_key: usize) -> Option> { 180 | let queries = &mut *self.inner.0.lock().unwrap(); 181 | 182 | queries.nodes[item_key].value.result.take() 183 | } 184 | 185 | fn remove(&self, item_key: usize, registration: event::Registration) { 186 | let queries = &mut *self.inner.0.lock().unwrap(); 187 | 188 | // remove from next list if present 189 | queries.next.remove(&mut queries.nodes, item_key); 190 | 191 | let qi = queries.nodes.remove(item_key).value; 192 | 193 | if let Some(invalidated) = &qi.invalidated { 194 | invalidated.store(true, Ordering::Relaxed); 195 | } 196 | 197 | queries 198 | .registrations 199 | .push_back((registration, qi.set_readiness)); 200 | } 201 | 202 | #[cfg(test)] 203 | fn invalidated_count(&self) -> u32 { 204 | let queries = &mut *self.inner.0.lock().unwrap(); 205 | 206 | queries.invalidated_count 207 | } 208 | } 209 | 210 | struct ResolverInner { 211 | workers: Vec>, 212 | queries: Queries, 213 | } 214 | 215 | impl ResolverInner { 216 | fn new(num_threads: usize, queries_max: usize, resolve_fn: Arc) -> Self 217 | where 218 | F: Fn(&str) -> Result + Send + Sync + 'static, 219 | { 220 | let mut workers = Vec::with_capacity(num_threads); 221 | let queries = Queries::new(queries_max); 222 | 223 | for _ in 0..workers.capacity() { 224 | let queries = queries.clone(); 225 | let resolve_fn = resolve_fn.clone(); 226 | 227 | let thread = thread::Builder::new() 228 | .name("resolver".to_string()) 229 | .spawn(move || { 230 | let invalidated = Arc::new(AtomicBool::new(false)); 231 | 232 | loop { 233 | assert_eq!(Arc::strong_count(&invalidated), 1); 234 | 235 | let (item_key, host) = match queries.get_next(&invalidated) { 236 | Some(ret) => ret, 237 | None => break, 238 | }; 239 | 240 | let ret = resolve_fn(host.as_str()); 241 | 242 | queries.set_result(item_key, ret, &invalidated); 243 | } 244 | }) 245 | .unwrap(); 246 | 247 | workers.push(thread); 248 | } 249 | 250 | Self { workers, queries } 251 | } 252 | 253 | #[allow(clippy::result_unit_err)] 254 | fn resolve(&self, host: &str) -> Result { 255 | let (item_key, reg) = self.queries.add(host)?; 256 | 257 | Ok(Query { 258 | queries: self.queries.clone(), 259 | item_key, 260 | registration: Some(reg), 261 | }) 262 | } 263 | 264 | fn stop(&mut self) { 265 | self.queries.set_stop_flag(); 266 | 267 | for worker in self.workers.drain(..) { 268 | worker.join().unwrap(); 269 | } 270 | } 271 | } 272 | 273 | impl Drop for ResolverInner { 274 | fn drop(&mut self) { 275 | self.stop(); 276 | } 277 | } 278 | 279 | pub struct Resolver { 280 | inner: ResolverInner, 281 | } 282 | 283 | impl Resolver { 284 | pub fn new(num_threads: usize, queries_max: usize) -> Self { 285 | let inner = ResolverInner::new(num_threads, queries_max, Arc::new(std_resolve)); 286 | 287 | Self { inner } 288 | } 289 | 290 | #[allow(clippy::result_unit_err)] 291 | pub fn resolve(&self, host: &str) -> Result { 292 | self.inner.resolve(host) 293 | } 294 | } 295 | 296 | pub struct Query { 297 | queries: Queries, 298 | item_key: usize, 299 | registration: Option, 300 | } 301 | 302 | impl Query { 303 | pub fn get_read_registration(&self) -> &event::Registration { 304 | self.registration.as_ref().unwrap() 305 | } 306 | 307 | pub fn process(&self) -> Option> { 308 | self.queries.take_result(self.item_key) 309 | } 310 | } 311 | 312 | impl Drop for Query { 313 | fn drop(&mut self) { 314 | let reg = self.registration.take().unwrap(); 315 | 316 | self.queries.remove(self.item_key, reg); 317 | } 318 | } 319 | 320 | #[cfg(test)] 321 | mod tests { 322 | use super::*; 323 | 324 | #[test] 325 | fn resolve() { 326 | let mut poller = event::Poller::new(1).unwrap(); 327 | 328 | let resolver = Resolver::new(1, 1); 329 | 330 | let query = resolver.resolve("127.0.0.1").unwrap(); 331 | 332 | // queries_max is 1, so this should error 333 | assert_eq!(resolver.resolve("127.0.0.1").is_err(), true); 334 | 335 | // register query interest with poller 336 | poller 337 | .register_custom( 338 | query.get_read_registration(), 339 | mio::Token(1), 340 | Interest::READABLE, 341 | ) 342 | .unwrap(); 343 | 344 | // wait for completion 345 | let result = loop { 346 | if let Some(result) = query.process() { 347 | break result; 348 | } 349 | 350 | poller.poll(None).unwrap(); 351 | 352 | for _ in poller.iter_events() {} 353 | }; 354 | 355 | // deregister query interest 356 | poller 357 | .deregister_custom(query.get_read_registration()) 358 | .unwrap(); 359 | 360 | assert_eq!(result.unwrap().as_slice(), &[IpAddr::from([127, 0, 0, 1])]); 361 | } 362 | 363 | #[test] 364 | fn invalidate_query() { 365 | let mut inner = { 366 | let cond = Arc::new((Mutex::new(false), Condvar::new())); 367 | 368 | let resolve_fn = { 369 | let cond = cond.clone(); 370 | 371 | Arc::new(move |_: &str| { 372 | let (lock, cvar) = &*cond; 373 | 374 | let guard = lock.lock().unwrap(); 375 | 376 | // let main thread know we've started 377 | cvar.notify_one(); 378 | 379 | // wait for query to be removed 380 | let _guard = cvar.wait(guard).unwrap(); 381 | 382 | Ok(Addrs::new()) 383 | }) 384 | }; 385 | 386 | let (lock, cvar) = &*cond; 387 | let guard = lock.lock().unwrap(); 388 | 389 | let inner = ResolverInner::new(1, 1, resolve_fn); 390 | 391 | let query = inner.resolve("127.0.0.1").unwrap(); 392 | 393 | // wait for resolve_fn to start 394 | let _guard = cvar.wait(guard).unwrap(); 395 | 396 | drop(query); 397 | 398 | // let worker know the query has been removed 399 | cvar.notify_one(); 400 | 401 | inner 402 | }; 403 | 404 | inner.stop(); 405 | 406 | assert_eq!(inner.queries.invalidated_count(), 1); 407 | } 408 | } 409 | -------------------------------------------------------------------------------- /src/app.rs: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2020-2023 Fanout, Inc. 3 | * 4 | * Licensed under the Apache License, Version 2.0 (the "License"); 5 | * you may not use this file except in compliance with the License. 6 | * You may obtain a copy of the License at 7 | * 8 | * http://www.apache.org/licenses/LICENSE-2.0 9 | * 10 | * Unless required by applicable law or agreed to in writing, software 11 | * distributed under the License is distributed on an "AS IS" BASIS, 12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | * See the License for the specific language governing permissions and 14 | * limitations under the License. 15 | */ 16 | 17 | use crate::client::Client; 18 | use crate::server::{Server, MSG_RETAINED_PER_CONNECTION_MAX, MSG_RETAINED_PER_WORKER_MAX}; 19 | use crate::websocket; 20 | use crate::zhttpsocket; 21 | use crate::zmq::SpecInfo; 22 | use ipnet::IpNet; 23 | use log::info; 24 | use signal_hook; 25 | use signal_hook::consts::TERM_SIGNALS; 26 | use signal_hook::iterator::Signals; 27 | use std::cmp; 28 | use std::path::PathBuf; 29 | use std::sync::atomic::AtomicBool; 30 | use std::sync::Arc; 31 | use std::time::Duration; 32 | 33 | const INIT_HWM: usize = 128; 34 | 35 | fn make_specs(base: &str, is_server: bool) -> Result<(String, String, String), String> { 36 | if base.starts_with("ipc:") { 37 | if is_server { 38 | Ok(( 39 | format!("{}-{}", base, "in"), 40 | format!("{}-{}", base, "in-stream"), 41 | format!("{}-{}", base, "out"), 42 | )) 43 | } else { 44 | Ok(( 45 | format!("{}-{}", base, "out"), 46 | format!("{}-{}", base, "out-stream"), 47 | format!("{}-{}", base, "in"), 48 | )) 49 | } 50 | } else if base.starts_with("tcp:") { 51 | match base.rfind(':') { 52 | Some(pos) => match base[(pos + 1)..base.len()].parse::() { 53 | Ok(port) => Ok(( 54 | format!("{}:{}", &base[..pos], port), 55 | format!("{}:{}", &base[..pos], port + 1), 56 | format!("{}:{}", &base[..pos], port + 2), 57 | )), 58 | Err(e) => Err(format!("error parsing tcp port in base spec: {}", e)), 59 | }, 60 | None => Err("tcp base spec must specify port".into()), 61 | } 62 | } else { 63 | Err("base spec must be ipc or tcp".into()) 64 | } 65 | } 66 | 67 | pub enum ListenSpec { 68 | Tcp { 69 | addr: std::net::SocketAddr, 70 | tls: bool, 71 | default_cert: Option, 72 | }, 73 | Local { 74 | path: PathBuf, 75 | mode: Option, 76 | user: Option, 77 | group: Option, 78 | }, 79 | } 80 | 81 | pub struct ListenConfig { 82 | pub spec: ListenSpec, 83 | pub stream: bool, 84 | } 85 | 86 | pub struct Config { 87 | pub instance_id: String, 88 | pub workers: usize, 89 | pub req_maxconn: usize, 90 | pub stream_maxconn: usize, 91 | pub buffer_size: usize, 92 | pub body_buffer_size: usize, 93 | pub messages_max: usize, 94 | pub req_timeout: Duration, 95 | pub stream_timeout: Duration, 96 | pub listen: Vec, 97 | pub zclient_req: Vec, 98 | pub zclient_stream: Vec, 99 | pub zclient_connect: bool, 100 | pub zserver_req: Vec, 101 | pub zserver_stream: Vec, 102 | pub zserver_connect: bool, 103 | pub ipc_file_mode: u32, 104 | pub certs_dir: PathBuf, 105 | pub allow_compression: bool, 106 | pub deny: Vec, 107 | } 108 | 109 | pub struct App { 110 | _server: Option, 111 | _client: Option, 112 | } 113 | 114 | impl App { 115 | pub fn new(config: &Config) -> Result { 116 | if config.req_maxconn < config.workers { 117 | return Err("req maxconn must be >= workers".into()); 118 | } 119 | 120 | if config.stream_maxconn < config.workers { 121 | return Err("stream maxconn must be >= workers".into()); 122 | } 123 | 124 | let zmq_context = Arc::new(zmq::Context::new()); 125 | 126 | // set hwm to 5% of maxconn 127 | let other_hwm = cmp::max((config.req_maxconn + config.stream_maxconn) / 20, 1); 128 | 129 | let handle_bound = cmp::max(other_hwm / config.workers, 1); 130 | 131 | let maxconn = config.req_maxconn + config.stream_maxconn; 132 | 133 | let server = if !config.listen.is_empty() { 134 | let mut any_req = false; 135 | let mut any_stream = false; 136 | 137 | for lc in config.listen.iter() { 138 | if lc.stream { 139 | any_stream = true; 140 | } else { 141 | any_req = true; 142 | } 143 | } 144 | 145 | let mut zsockman = zhttpsocket::ClientSocketManager::new( 146 | Arc::clone(&zmq_context), 147 | &config.instance_id, 148 | (MSG_RETAINED_PER_CONNECTION_MAX * maxconn) 149 | + (MSG_RETAINED_PER_WORKER_MAX * config.workers), 150 | INIT_HWM, 151 | other_hwm, 152 | handle_bound, 153 | ); 154 | 155 | if any_req { 156 | let mut specs = Vec::new(); 157 | 158 | for spec in config.zclient_req.iter() { 159 | if config.zclient_connect { 160 | info!("zhttp client connect {}", spec); 161 | } else { 162 | info!("zhttp client bind {}", spec); 163 | } 164 | 165 | specs.push(SpecInfo { 166 | spec: spec.clone(), 167 | bind: !config.zclient_connect, 168 | ipc_file_mode: config.ipc_file_mode, 169 | }); 170 | } 171 | 172 | if let Err(e) = zsockman.set_client_req_specs(&specs) { 173 | return Err(format!("failed to set zhttp client req specs: {}", e)); 174 | } 175 | } 176 | 177 | if any_stream { 178 | let mut out_specs = Vec::new(); 179 | let mut out_stream_specs = Vec::new(); 180 | let mut in_specs = Vec::new(); 181 | 182 | for spec in config.zclient_stream.iter() { 183 | let (out_spec, out_stream_spec, in_spec) = make_specs(spec, false)?; 184 | 185 | if config.zclient_connect { 186 | info!( 187 | "zhttp client connect {} {} {}", 188 | out_spec, out_stream_spec, in_spec 189 | ); 190 | } else { 191 | info!( 192 | "zhttp client bind {} {} {}", 193 | out_spec, out_stream_spec, in_spec 194 | ); 195 | } 196 | 197 | out_specs.push(SpecInfo { 198 | spec: out_spec, 199 | bind: !config.zclient_connect, 200 | ipc_file_mode: config.ipc_file_mode, 201 | }); 202 | 203 | out_stream_specs.push(SpecInfo { 204 | spec: out_stream_spec, 205 | bind: !config.zclient_connect, 206 | ipc_file_mode: config.ipc_file_mode, 207 | }); 208 | 209 | in_specs.push(SpecInfo { 210 | spec: in_spec, 211 | bind: !config.zclient_connect, 212 | ipc_file_mode: config.ipc_file_mode, 213 | }); 214 | } 215 | 216 | if let Err(e) = 217 | zsockman.set_client_stream_specs(&out_specs, &out_stream_specs, &in_specs) 218 | { 219 | return Err(format!("failed to set zhttp client stream specs: {}", e)); 220 | } 221 | } 222 | 223 | Some(Server::new( 224 | &config.instance_id, 225 | config.workers, 226 | config.req_maxconn, 227 | config.stream_maxconn, 228 | config.buffer_size, 229 | config.body_buffer_size, 230 | config.messages_max, 231 | config.req_timeout, 232 | config.stream_timeout, 233 | &config.listen, 234 | config.certs_dir.as_path(), 235 | config.allow_compression, 236 | zsockman, 237 | handle_bound, 238 | )?) 239 | } else { 240 | None 241 | }; 242 | 243 | let client = if !config.zserver_req.is_empty() || !config.zserver_stream.is_empty() { 244 | let mut zsockman = zhttpsocket::ServerSocketManager::new( 245 | Arc::clone(&zmq_context), 246 | &config.instance_id, 247 | (MSG_RETAINED_PER_CONNECTION_MAX * maxconn) 248 | + (MSG_RETAINED_PER_WORKER_MAX * config.workers), 249 | INIT_HWM, 250 | other_hwm, 251 | handle_bound, 252 | config.stream_maxconn, 253 | ); 254 | 255 | if !config.zserver_req.is_empty() { 256 | let mut specs = Vec::new(); 257 | 258 | for spec in config.zserver_req.iter() { 259 | if config.zserver_connect { 260 | info!("zhttp server connect {}", spec); 261 | } else { 262 | info!("zhttp server bind {}", spec); 263 | } 264 | 265 | specs.push(SpecInfo { 266 | spec: spec.clone(), 267 | bind: !config.zserver_connect, 268 | ipc_file_mode: config.ipc_file_mode, 269 | }); 270 | } 271 | 272 | if let Err(e) = zsockman.set_server_req_specs(&specs) { 273 | return Err(format!("failed to set zhttp server req specs: {}", e)); 274 | } 275 | } 276 | 277 | let zsockman = Arc::new(zsockman); 278 | 279 | let client = Client::new( 280 | &config.instance_id, 281 | config.workers, 282 | config.req_maxconn, 283 | config.stream_maxconn, 284 | config.buffer_size, 285 | config.body_buffer_size, 286 | config.messages_max, 287 | config.req_timeout, 288 | config.stream_timeout, 289 | config.allow_compression, 290 | &config.deny, 291 | zsockman.clone(), 292 | handle_bound, 293 | )?; 294 | 295 | // stream specs must only be applied after client is initialized 296 | if !config.zserver_stream.is_empty() { 297 | let mut in_specs = Vec::new(); 298 | let mut in_stream_specs = Vec::new(); 299 | let mut out_specs = Vec::new(); 300 | 301 | for spec in config.zserver_stream.iter() { 302 | let (in_spec, in_stream_spec, out_spec) = make_specs(spec, true)?; 303 | 304 | if config.zserver_connect { 305 | info!( 306 | "zhttp server connect {} {} {}", 307 | in_spec, in_stream_spec, out_spec 308 | ); 309 | } else { 310 | info!( 311 | "zhttp server bind {} {} {}", 312 | in_spec, in_stream_spec, out_spec 313 | ); 314 | } 315 | 316 | in_specs.push(SpecInfo { 317 | spec: in_spec, 318 | bind: !config.zserver_connect, 319 | ipc_file_mode: config.ipc_file_mode, 320 | }); 321 | 322 | in_stream_specs.push(SpecInfo { 323 | spec: in_stream_spec, 324 | bind: !config.zserver_connect, 325 | ipc_file_mode: config.ipc_file_mode, 326 | }); 327 | 328 | out_specs.push(SpecInfo { 329 | spec: out_spec, 330 | bind: !config.zserver_connect, 331 | ipc_file_mode: config.ipc_file_mode, 332 | }); 333 | } 334 | 335 | if let Err(e) = 336 | zsockman.set_server_stream_specs(&in_specs, &in_stream_specs, &out_specs) 337 | { 338 | return Err(format!("failed to set zhttp server stream specs: {}", e)); 339 | } 340 | } 341 | 342 | Some(client) 343 | } else { 344 | None 345 | }; 346 | 347 | Ok(Self { 348 | _server: server, 349 | _client: client, 350 | }) 351 | } 352 | 353 | pub fn wait_for_term(&self) { 354 | let mut signals = Signals::new(TERM_SIGNALS).unwrap(); 355 | 356 | let term_now = Arc::new(AtomicBool::new(false)); 357 | 358 | // ensure two term signals in a row causes the app to immediately exit 359 | for signal_type in TERM_SIGNALS { 360 | signal_hook::flag::register_conditional_shutdown( 361 | *signal_type, 362 | 1, // exit code 363 | Arc::clone(&term_now), 364 | ) 365 | .unwrap(); 366 | 367 | signal_hook::flag::register(*signal_type, Arc::clone(&term_now)).unwrap(); 368 | } 369 | 370 | // wait for termination 371 | for signal in &mut signals { 372 | match signal { 373 | signal_type if TERM_SIGNALS.contains(&signal_type) => break, 374 | _ => unreachable!(), 375 | } 376 | } 377 | } 378 | 379 | pub fn sizes() -> Vec<(String, usize)> { 380 | let mut out = Vec::new(); 381 | 382 | out.extend(Server::task_sizes()); 383 | out.extend(Client::task_sizes()); 384 | out.push(( 385 | "deflate_codec_state".to_string(), 386 | websocket::deflate_codec_state_size(), 387 | )); 388 | 389 | out 390 | } 391 | } 392 | -------------------------------------------------------------------------------- /src/zmq.rs: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2020-2023 Fanout, Inc. 3 | * 4 | * Licensed under the Apache License, Version 2.0 (the "License"); 5 | * you may not use this file except in compliance with the License. 6 | * You may obtain a copy of the License at 7 | * 8 | * http://www.apache.org/licenses/LICENSE-2.0 9 | * 10 | * Unless required by applicable law or agreed to in writing, software 11 | * distributed under the License is distributed on an "AS IS" BASIS, 12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | * See the License for the specific language governing permissions and 14 | * limitations under the License. 15 | */ 16 | 17 | use arrayvec::ArrayVec; 18 | use std::cell::Cell; 19 | use std::cell::RefCell; 20 | use std::fmt; 21 | use std::fs; 22 | use std::io; 23 | use std::os::unix::fs::PermissionsExt; 24 | 25 | const MULTIPART_HEADERS_MAX: usize = 8; 26 | 27 | fn trim_prefix<'a>(s: &'a str, prefix: &str) -> Result<&'a str, ()> { 28 | if let Some(s) = s.strip_prefix(prefix) { 29 | Ok(s) 30 | } else { 31 | Err(()) 32 | } 33 | } 34 | 35 | #[derive(Clone)] 36 | pub struct SpecInfo { 37 | pub spec: String, 38 | pub bind: bool, 39 | pub ipc_file_mode: u32, 40 | } 41 | 42 | impl fmt::Display for SpecInfo { 43 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 44 | if self.bind { 45 | write!(f, "bind:{}", self.spec) 46 | } else { 47 | write!(f, "connect:{}", self.spec) 48 | } 49 | } 50 | } 51 | 52 | impl fmt::Debug for SpecInfo { 53 | fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { 54 | write!(f, "{}", self) 55 | } 56 | } 57 | 58 | #[derive(Debug)] 59 | pub enum ZmqSocketError { 60 | Connect(String, zmq::Error), 61 | Bind(String, zmq::Error), 62 | SetMode(String, io::Error), 63 | } 64 | 65 | impl ToString for ZmqSocketError { 66 | fn to_string(&self) -> String { 67 | match self { 68 | ZmqSocketError::Connect(spec, e) => format!("connect {}: {}", spec, e), 69 | ZmqSocketError::Bind(spec, e) => format!("bind {}: {}", spec, e), 70 | ZmqSocketError::SetMode(spec, e) => format!("set mode {}: {}", spec, e), 71 | } 72 | } 73 | } 74 | 75 | #[derive(Clone)] 76 | struct ActiveSpec { 77 | pub spec: SpecInfo, 78 | pub endpoint: String, 79 | } 80 | 81 | fn unbind(sock: &zmq::Socket, endpoint: &str) -> zmq::Result<()> { 82 | // NOTE: use zmq_unbind instead when it becomes available in rust-zmq 83 | sock.disconnect(endpoint) 84 | } 85 | 86 | fn setup_spec(sock: &zmq::Socket, spec: &SpecInfo) -> Result { 87 | if spec.bind { 88 | match sock.bind(&spec.spec) { 89 | Ok(_) => { 90 | let endpoint = sock.get_last_endpoint().unwrap().unwrap(); 91 | 92 | if let Ok(path) = trim_prefix(&spec.spec, "ipc://") { 93 | if spec.ipc_file_mode > 0 { 94 | let perms = fs::Permissions::from_mode(spec.ipc_file_mode); 95 | if let Err(e) = fs::set_permissions(path, perms) { 96 | // if setting perms fails, undo the bind 97 | unbind(sock, &endpoint).unwrap(); 98 | 99 | return Err(ZmqSocketError::SetMode(spec.spec.clone(), e)); 100 | } 101 | } 102 | } 103 | 104 | Ok(endpoint) 105 | } 106 | Err(e) => Err(ZmqSocketError::Bind(spec.spec.clone(), e)), 107 | } 108 | } else { 109 | match sock.connect(&spec.spec) { 110 | Ok(_) => Ok(spec.spec.clone()), 111 | Err(e) => Err(ZmqSocketError::Connect(spec.spec.clone(), e)), 112 | } 113 | } 114 | } 115 | 116 | fn unsetup_spec(sock: &zmq::Socket, spec: &ActiveSpec) { 117 | if spec.spec.bind { 118 | unbind(sock, &spec.endpoint).unwrap(); 119 | 120 | if let Ok(path) = trim_prefix(&spec.endpoint, "ipc://") { 121 | if fs::remove_file(path).is_err() { 122 | // oh well, we tried 123 | } 124 | } 125 | } else { 126 | sock.disconnect(&spec.endpoint).unwrap(); 127 | } 128 | } 129 | 130 | pub type MultipartHeader = ArrayVec; 131 | 132 | pub struct ZmqSocket { 133 | inner: zmq::Socket, 134 | events: Cell, 135 | specs: RefCell>, 136 | } 137 | 138 | impl ZmqSocket { 139 | pub fn new(ctx: &zmq::Context, socket_type: zmq::SocketType) -> Self { 140 | Self { 141 | inner: ctx.socket(socket_type).unwrap(), 142 | events: Cell::new(zmq::PollEvents::empty()), 143 | specs: RefCell::new(Vec::new()), 144 | } 145 | } 146 | 147 | pub fn inner(&self) -> &zmq::Socket { 148 | &self.inner 149 | } 150 | 151 | pub fn update_events(&self) { 152 | loop { 153 | match self.inner.get_events() { 154 | Ok(events) => { 155 | self.events.set(events); 156 | break; 157 | } 158 | Err(zmq::Error::EINTR) => continue, 159 | Err(e) => panic!("get events error: {}", e), 160 | } 161 | } 162 | } 163 | 164 | pub fn events(&self) -> zmq::PollEvents { 165 | self.events.get() 166 | } 167 | 168 | pub fn send(&self, msg: zmq::Message, flags: i32) -> Result<(), zmq::Error> { 169 | let flags = flags & zmq::DONTWAIT; 170 | 171 | if let Err(e) = self.inner.send(msg, flags) { 172 | self.update_events(); 173 | return Err(e); 174 | } 175 | 176 | self.update_events(); 177 | 178 | Ok(()) 179 | } 180 | 181 | pub fn send_to( 182 | &self, 183 | header: &MultipartHeader, 184 | content: zmq::Message, 185 | flags: i32, 186 | ) -> Result<(), zmq::Error> { 187 | let mut headers: ArrayVec<&[u8], MULTIPART_HEADERS_MAX> = ArrayVec::new(); 188 | 189 | for part in header { 190 | headers.push(part); 191 | } 192 | 193 | let flags = flags & zmq::DONTWAIT; 194 | 195 | if let Err(e) = self.inner.send_multipart(&headers, flags | zmq::SNDMORE) { 196 | self.update_events(); 197 | return Err(e); 198 | } 199 | 200 | if let Err(e) = self.inner.send(zmq::Message::new(), flags | zmq::SNDMORE) { 201 | self.update_events(); 202 | return Err(e); 203 | } 204 | 205 | self.send(content, flags) 206 | } 207 | 208 | pub fn recv(&self, flags: i32) -> Result { 209 | let flags = flags & zmq::DONTWAIT; 210 | 211 | // get the first part 212 | let msg = match self.inner.recv_msg(flags) { 213 | Ok(msg) => msg, 214 | Err(e) => { 215 | self.update_events(); 216 | return Err(e); 217 | } 218 | }; 219 | 220 | let flags = 0; 221 | 222 | // eat the rest of the parts 223 | while self.inner.get_rcvmore().unwrap() { 224 | self.inner.recv_msg(flags).unwrap(); 225 | } 226 | 227 | self.update_events(); 228 | 229 | Ok(msg) 230 | } 231 | 232 | pub fn recv_routed(&self, flags: i32) -> Result<(MultipartHeader, zmq::Message), zmq::Error> { 233 | let flags = flags & zmq::DONTWAIT; 234 | 235 | let mut header = MultipartHeader::new(); 236 | 237 | loop { 238 | // read parts until we reach the separator 239 | match self.inner.recv_msg(flags) { 240 | Ok(msg) => { 241 | if msg.is_empty() { 242 | break; 243 | } 244 | 245 | if header.len() == header.capacity() { 246 | // header too large 247 | 248 | let flags = 0; 249 | 250 | // eat the rest of the parts 251 | while self.inner.get_rcvmore().unwrap() { 252 | self.inner.recv_msg(flags).unwrap(); 253 | } 254 | 255 | self.update_events(); 256 | 257 | return Err(zmq::Error::EINVAL); 258 | } 259 | 260 | header.push(msg); 261 | } 262 | Err(e) => { 263 | self.update_events(); 264 | return Err(e); 265 | } 266 | } 267 | } 268 | 269 | let flags = 0; 270 | 271 | // if we get here, we've read the separator. content parts should follow 272 | 273 | if !self.inner.get_rcvmore().unwrap() { 274 | return Err(zmq::Error::EINVAL); 275 | } 276 | 277 | // get the first part of the content 278 | let msg = match self.inner.recv_msg(flags) { 279 | Ok(msg) => msg, 280 | Err(e) => { 281 | self.update_events(); 282 | return Err(e); 283 | } 284 | }; 285 | 286 | // eat the rest of the parts 287 | while self.inner.get_rcvmore().unwrap() { 288 | self.inner.recv_msg(flags).unwrap(); 289 | } 290 | 291 | self.update_events(); 292 | 293 | Ok((header, msg)) 294 | } 295 | 296 | pub fn apply_specs(&self, new_specs: &[SpecInfo]) -> Result<(), ZmqSocketError> { 297 | let mut specs = self.specs.borrow_mut(); 298 | 299 | let mut to_remove = Vec::new(); 300 | for cur in specs.iter() { 301 | let mut found = false; 302 | for new in new_specs.iter() { 303 | if cur.spec.spec == new.spec && cur.spec.bind == new.bind { 304 | found = true; 305 | break; 306 | } 307 | } 308 | if !found { 309 | to_remove.push(cur.clone()); 310 | } 311 | } 312 | 313 | let mut to_add = Vec::new(); 314 | let mut to_update = Vec::new(); 315 | for new in new_specs.iter() { 316 | let mut found = None; 317 | for (ci, cur) in specs.iter().enumerate() { 318 | if new.spec == cur.spec.spec && new.bind == cur.spec.bind { 319 | found = Some(ci); 320 | break; 321 | } 322 | } 323 | match found { 324 | Some(ci) => { 325 | if new.ipc_file_mode != specs[ci].spec.ipc_file_mode { 326 | to_update.push(new.clone()); 327 | } 328 | } 329 | None => { 330 | to_add.push(new.clone()); 331 | } 332 | } 333 | } 334 | 335 | let mut added = Vec::new(); 336 | 337 | // add specs we dont have. on fail, undo them 338 | for spec in to_add.iter() { 339 | match setup_spec(&self.inner, spec) { 340 | Ok(endpoint) => { 341 | added.push(ActiveSpec { 342 | spec: spec.clone(), 343 | endpoint, 344 | }); 345 | } 346 | Err(e) => { 347 | // undo previous adds 348 | for spec in added.iter().rev() { 349 | unsetup_spec(&self.inner, spec); 350 | } 351 | return Err(e); 352 | } 353 | } 354 | } 355 | 356 | // update ipc file mode 357 | let mut prev_perms = Vec::new(); 358 | for spec in to_update.iter() { 359 | let mut err = None; 360 | 361 | if let Ok(path) = trim_prefix(&spec.spec, "ipc://") { 362 | if spec.ipc_file_mode > 0 { 363 | match fs::metadata(path) { 364 | Ok(meta) => { 365 | let perms = fs::Permissions::from_mode(spec.ipc_file_mode); 366 | match fs::set_permissions(path, perms) { 367 | Ok(_) => { 368 | prev_perms.push((String::from(path), meta.permissions())); 369 | } 370 | Err(e) => { 371 | err = Some(ZmqSocketError::SetMode(spec.spec.clone(), e)); 372 | } 373 | } 374 | } 375 | Err(e) => { 376 | err = Some(ZmqSocketError::SetMode(spec.spec.clone(), e)); 377 | } 378 | } 379 | } 380 | } 381 | 382 | if let Some(err) = err { 383 | // undo previous perms changes 384 | for (path, perms) in prev_perms { 385 | if fs::set_permissions(path, perms).is_err() { 386 | // oh well, we tried 387 | } 388 | } 389 | 390 | // undo previous adds 391 | for spec in added.iter().rev() { 392 | unsetup_spec(&self.inner, spec); 393 | } 394 | 395 | return Err(err); 396 | } 397 | } 398 | 399 | for spec in to_remove.iter() { 400 | unsetup_spec(&self.inner, spec); 401 | } 402 | 403 | // move current specs aside 404 | let prev_specs = std::mem::take(&mut *specs); 405 | 406 | // recompute current specs 407 | for new in new_specs { 408 | let mut s = None; 409 | 410 | // is it one we added? 411 | for spec in added.iter() { 412 | if new.spec == spec.spec.spec && new.bind == spec.spec.bind { 413 | s = Some(spec.clone()); 414 | break; 415 | } 416 | } 417 | 418 | // else, it must be one we had already 419 | if s.is_none() { 420 | for spec in prev_specs.iter() { 421 | if new.spec == spec.spec.spec && new.bind == spec.spec.bind { 422 | s = Some(spec.clone()); 423 | break; 424 | } 425 | } 426 | } 427 | 428 | assert!(s.is_some()); 429 | 430 | specs.push(s.unwrap()); 431 | } 432 | 433 | Ok(()) 434 | } 435 | } 436 | 437 | impl Drop for ZmqSocket { 438 | fn drop(&mut self) { 439 | let specs = self.specs.borrow(); 440 | 441 | for spec in specs.iter() { 442 | unsetup_spec(&self.inner, spec); 443 | } 444 | } 445 | } 446 | 447 | #[cfg(test)] 448 | mod tests { 449 | use super::*; 450 | 451 | #[test] 452 | fn test_send_after_disconnect() { 453 | let zmq_context = zmq::Context::new(); 454 | 455 | let s = ZmqSocket::new(&zmq_context, zmq::REQ); 456 | s.apply_specs(&[SpecInfo { 457 | spec: String::from("inproc://send-test"), 458 | bind: true, 459 | ipc_file_mode: 0, 460 | }]) 461 | .unwrap(); 462 | 463 | assert_eq!(s.events().contains(zmq::POLLOUT), false); 464 | 465 | let r = ZmqSocket::new(&zmq_context, zmq::REP); 466 | r.apply_specs(&[SpecInfo { 467 | spec: String::from("inproc://send-test"), 468 | bind: false, 469 | ipc_file_mode: 0, 470 | }]) 471 | .unwrap(); 472 | 473 | s.update_events(); 474 | 475 | assert_eq!(s.events().contains(zmq::POLLOUT), true); 476 | 477 | drop(r); 478 | 479 | assert_eq!( 480 | s.send((&b"test"[..]).into(), zmq::DONTWAIT), 481 | Err(zmq::Error::EAGAIN) 482 | ); 483 | 484 | assert_eq!(s.events().contains(zmq::POLLOUT), false); 485 | } 486 | } 487 | -------------------------------------------------------------------------------- /src/executor.rs: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2020-2023 Fanout, Inc. 3 | * 4 | * Licensed under the Apache License, Version 2.0 (the "License"); 5 | * you may not use this file except in compliance with the License. 6 | * You may obtain a copy of the License at 7 | * 8 | * http://www.apache.org/licenses/LICENSE-2.0 9 | * 10 | * Unless required by applicable law or agreed to in writing, software 11 | * distributed under the License is distributed on an "AS IS" BASIS, 12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | * See the License for the specific language governing permissions and 14 | * limitations under the License. 15 | */ 16 | 17 | use crate::list; 18 | use crate::waker; 19 | use log::debug; 20 | use slab::Slab; 21 | use std::cell::RefCell; 22 | use std::future::Future; 23 | use std::io; 24 | use std::mem; 25 | use std::pin::Pin; 26 | use std::rc::{Rc, Weak}; 27 | use std::task::{Context, Waker}; 28 | use std::time::Duration; 29 | 30 | thread_local! { 31 | static EXECUTOR: RefCell>> = RefCell::new(None); 32 | } 33 | 34 | type BoxFuture = Pin>>; 35 | 36 | struct TaskWaker { 37 | tasks: Weak, 38 | task_id: usize, 39 | } 40 | 41 | impl waker::RcWake for TaskWaker { 42 | fn wake(self: Rc) { 43 | if let Some(tasks) = self.tasks.upgrade() { 44 | tasks.wake(self.task_id); 45 | } 46 | } 47 | } 48 | 49 | fn poll_fut(fut: &mut BoxFuture, waker: Waker) -> bool { 50 | // convert from Pin to Pin<&mut> 51 | let fut: Pin<&mut dyn Future> = fut.as_mut(); 52 | 53 | let mut cx = Context::from_waker(&waker); 54 | 55 | fut.poll(&mut cx).is_ready() 56 | } 57 | 58 | struct Task { 59 | fut: Option>>>, 60 | wakeable: bool, 61 | } 62 | 63 | struct TasksData { 64 | nodes: Slab>, 65 | next: list::List, 66 | wakers: Vec>, 67 | } 68 | 69 | struct Tasks { 70 | data: RefCell, 71 | pre_poll: RefCell>>, 72 | } 73 | 74 | impl Tasks { 75 | fn new(max: usize) -> Rc { 76 | let data = TasksData { 77 | nodes: Slab::with_capacity(max), 78 | next: list::List::default(), 79 | wakers: Vec::with_capacity(max), 80 | }; 81 | 82 | let tasks = Rc::new(Self { 83 | data: RefCell::new(data), 84 | pre_poll: RefCell::new(None), 85 | }); 86 | 87 | { 88 | let data = &mut *tasks.data.borrow_mut(); 89 | 90 | for task_id in 0..data.nodes.capacity() { 91 | data.wakers.push(Rc::new(TaskWaker { 92 | tasks: Rc::downgrade(&tasks), 93 | task_id, 94 | })); 95 | } 96 | } 97 | 98 | tasks 99 | } 100 | 101 | fn is_empty(&self) -> bool { 102 | self.data.borrow().nodes.is_empty() 103 | } 104 | 105 | fn have_next(&self) -> bool { 106 | !self.data.borrow().next.is_empty() 107 | } 108 | 109 | fn add(&self, fut: F) -> Result<(), ()> 110 | where 111 | F: Future + 'static, 112 | { 113 | let data = &mut *self.data.borrow_mut(); 114 | 115 | if data.nodes.len() == data.nodes.capacity() { 116 | return Err(()); 117 | } 118 | 119 | let entry = data.nodes.vacant_entry(); 120 | let nkey = entry.key(); 121 | 122 | let task = Task { 123 | fut: Some(Box::pin(fut)), 124 | wakeable: false, 125 | }; 126 | 127 | entry.insert(list::Node::new(task)); 128 | 129 | data.next.push_back(&mut data.nodes, nkey); 130 | 131 | Ok(()) 132 | } 133 | 134 | fn remove(&self, task_id: usize) { 135 | let nkey = task_id; 136 | 137 | let data = &mut *self.data.borrow_mut(); 138 | 139 | let task = &mut data.nodes[nkey].value; 140 | 141 | // drop the future. this should cause it to drop any owned wakers 142 | task.fut = None; 143 | 144 | // at this point, we should be the only remaining owner 145 | assert_eq!(Rc::strong_count(&data.wakers[nkey]), 1); 146 | 147 | data.next.remove(&mut data.nodes, nkey); 148 | data.nodes.remove(nkey); 149 | } 150 | 151 | fn take_next_list(&self) -> list::List { 152 | let data = &mut *self.data.borrow_mut(); 153 | 154 | let mut l = list::List::default(); 155 | l.concat(&mut data.nodes, &mut data.next); 156 | 157 | l 158 | } 159 | 160 | fn append_to_next_list(&self, mut l: list::List) { 161 | let data = &mut *self.data.borrow_mut(); 162 | 163 | data.next.concat(&mut data.nodes, &mut l); 164 | } 165 | 166 | fn take_task(&self, l: &mut list::List) -> Option<(usize, BoxFuture, Waker)> { 167 | let nkey = match l.head { 168 | Some(nkey) => nkey, 169 | None => return None, 170 | }; 171 | 172 | let data = &mut *self.data.borrow_mut(); 173 | 174 | l.remove(&mut data.nodes, nkey); 175 | 176 | let task = &mut data.nodes[nkey].value; 177 | 178 | // both of these are cheap 179 | let fut = task.fut.take().unwrap(); 180 | let waker = waker::into_std(data.wakers[nkey].clone()); 181 | 182 | task.wakeable = true; 183 | 184 | Some((nkey, fut, waker)) 185 | } 186 | 187 | fn process_next(&self) { 188 | let mut l = self.take_next_list(); 189 | 190 | while let Some((task_id, mut fut, waker)) = self.take_task(&mut l) { 191 | self.pre_poll(); 192 | 193 | let done = poll_fut(&mut fut, waker); 194 | 195 | // take_task() took the future out of the task, so we 196 | // could poll it without having to maintain a borrow of 197 | // the tasks set. we'll put it back now 198 | self.set_fut(task_id, fut); 199 | 200 | if done { 201 | self.remove(task_id); 202 | } 203 | } 204 | } 205 | 206 | fn set_fut(&self, task_id: usize, fut: BoxFuture) { 207 | let nkey = task_id; 208 | 209 | let data = &mut *self.data.borrow_mut(); 210 | 211 | let task = &mut data.nodes[nkey].value; 212 | 213 | task.fut = Some(fut); 214 | } 215 | 216 | fn wake(&self, task_id: usize) { 217 | let nkey = task_id; 218 | 219 | let data = &mut *self.data.borrow_mut(); 220 | 221 | let node = &mut data.nodes[nkey]; 222 | 223 | if !node.value.wakeable { 224 | return; 225 | } 226 | 227 | node.value.wakeable = false; 228 | 229 | data.next.push_back(&mut data.nodes, nkey); 230 | } 231 | 232 | fn set_pre_poll(&self, pre_poll_fn: F) 233 | where 234 | F: FnMut() + 'static, 235 | { 236 | *self.pre_poll.borrow_mut() = Some(Box::new(pre_poll_fn)); 237 | } 238 | 239 | fn pre_poll(&self) { 240 | let pre_poll = &mut *self.pre_poll.borrow_mut(); 241 | 242 | if let Some(f) = pre_poll { 243 | f(); 244 | } 245 | } 246 | } 247 | 248 | pub struct Executor { 249 | tasks: Rc, 250 | } 251 | 252 | impl Executor { 253 | pub fn new(tasks_max: usize) -> Self { 254 | let tasks = Tasks::new(tasks_max); 255 | 256 | EXECUTOR.with(|ex| { 257 | if ex.borrow().is_some() { 258 | panic!("thread already has an Executor"); 259 | } 260 | 261 | ex.replace(Some(Rc::downgrade(&tasks))); 262 | }); 263 | 264 | Self { tasks } 265 | } 266 | 267 | #[allow(clippy::result_unit_err)] 268 | pub fn spawn(&self, fut: F) -> Result<(), ()> 269 | where 270 | F: Future + 'static, 271 | { 272 | debug!("spawning future with size {}", mem::size_of::()); 273 | 274 | self.tasks.add(fut) 275 | } 276 | 277 | pub fn set_pre_poll(&self, pre_poll_fn: F) 278 | where 279 | F: FnMut() + 'static, 280 | { 281 | self.tasks.set_pre_poll(pre_poll_fn); 282 | } 283 | 284 | pub fn have_tasks(&self) -> bool { 285 | !self.tasks.is_empty() 286 | } 287 | 288 | pub fn run_until_stalled(&self) { 289 | while self.tasks.have_next() { 290 | self.tasks.process_next() 291 | } 292 | } 293 | 294 | pub fn run(&self, mut park: F) -> Result<(), io::Error> 295 | where 296 | F: FnMut(Option) -> Result<(), io::Error>, 297 | { 298 | loop { 299 | self.tasks.process_next(); 300 | 301 | if !self.have_tasks() { 302 | break; 303 | } 304 | 305 | let (timeout, low_priority_tasks) = if self.tasks.have_next() { 306 | // some tasks trigger their own waker and return Pending in 307 | // order to achieve a yielding effect. in that case they will 308 | // already be queued up for processing again. move these 309 | // tasks aside so that they can be deprioritized, and use a 310 | // timeout of 0 when parking so we can quickly resume them 311 | 312 | let timeout = Duration::from_millis(0); 313 | let l = self.tasks.take_next_list(); 314 | 315 | (Some(timeout), Some(l)) 316 | } else { 317 | (None, None) 318 | }; 319 | 320 | park(timeout)?; 321 | 322 | // requeue any tasks that had yielded 323 | if let Some(l) = low_priority_tasks { 324 | self.tasks.append_to_next_list(l); 325 | } 326 | } 327 | 328 | Ok(()) 329 | } 330 | 331 | pub fn current() -> Option { 332 | EXECUTOR.with(|ex| { 333 | (*ex.borrow_mut()).as_mut().map(|tasks| Self { 334 | tasks: tasks.upgrade().unwrap(), 335 | }) 336 | }) 337 | } 338 | 339 | pub fn spawner(&self) -> Spawner { 340 | Spawner { 341 | tasks: Rc::downgrade(&self.tasks), 342 | } 343 | } 344 | } 345 | 346 | impl Drop for Executor { 347 | fn drop(&mut self) { 348 | EXECUTOR.with(|ex| { 349 | if Rc::strong_count(&self.tasks) == 1 { 350 | ex.replace(None); 351 | } 352 | }); 353 | } 354 | } 355 | 356 | pub struct Spawner { 357 | tasks: Weak, 358 | } 359 | 360 | impl Spawner { 361 | #[allow(clippy::result_unit_err)] 362 | pub fn spawn(&self, fut: F) -> Result<(), ()> 363 | where 364 | F: Future + 'static, 365 | { 366 | let tasks = match self.tasks.upgrade() { 367 | Some(tasks) => tasks, 368 | None => return Err(()), 369 | }; 370 | 371 | let ex = Executor { tasks }; 372 | 373 | ex.spawn(fut) 374 | } 375 | } 376 | 377 | #[cfg(test)] 378 | mod tests { 379 | use super::*; 380 | use std::cell::Cell; 381 | use std::mem; 382 | use std::task::Poll; 383 | 384 | struct TestFutureData { 385 | ready: bool, 386 | waker: Option, 387 | } 388 | 389 | struct TestFuture { 390 | data: Rc>, 391 | } 392 | 393 | impl TestFuture { 394 | fn new() -> Self { 395 | let data = TestFutureData { 396 | ready: false, 397 | waker: None, 398 | }; 399 | 400 | Self { 401 | data: Rc::new(RefCell::new(data)), 402 | } 403 | } 404 | 405 | fn handle(&self) -> TestHandle { 406 | TestHandle { 407 | data: Rc::clone(&self.data), 408 | } 409 | } 410 | } 411 | 412 | impl Future for TestFuture { 413 | type Output = (); 414 | 415 | fn poll(self: Pin<&mut Self>, cx: &mut Context) -> Poll { 416 | let mut data = self.data.borrow_mut(); 417 | 418 | match data.ready { 419 | true => Poll::Ready(()), 420 | false => { 421 | data.waker = Some(cx.waker().clone()); 422 | 423 | Poll::Pending 424 | } 425 | } 426 | } 427 | } 428 | 429 | struct TestHandle { 430 | data: Rc>, 431 | } 432 | 433 | impl TestHandle { 434 | fn set_ready(&self) { 435 | let data = &mut *self.data.borrow_mut(); 436 | 437 | data.ready = true; 438 | 439 | if let Some(waker) = data.waker.take() { 440 | waker.wake(); 441 | } 442 | } 443 | } 444 | 445 | struct EarlyWakeFuture { 446 | done: bool, 447 | } 448 | 449 | impl EarlyWakeFuture { 450 | fn new() -> Self { 451 | Self { done: false } 452 | } 453 | } 454 | 455 | impl Future for EarlyWakeFuture { 456 | type Output = (); 457 | 458 | fn poll(mut self: Pin<&mut Self>, cx: &mut Context) -> Poll { 459 | if !self.done { 460 | self.done = true; 461 | cx.waker().wake_by_ref(); 462 | 463 | return Poll::Pending; 464 | } 465 | 466 | Poll::Ready(()) 467 | } 468 | } 469 | 470 | #[test] 471 | fn test_executor_step() { 472 | let executor = Executor::new(1); 473 | 474 | let fut1 = TestFuture::new(); 475 | let fut2 = TestFuture::new(); 476 | 477 | let handle1 = fut1.handle(); 478 | let handle2 = fut2.handle(); 479 | 480 | let started = Rc::new(Cell::new(false)); 481 | let fut1_done = Rc::new(Cell::new(false)); 482 | let finishing = Rc::new(Cell::new(false)); 483 | 484 | { 485 | let started = Rc::clone(&started); 486 | let fut1_done = Rc::clone(&fut1_done); 487 | let finishing = Rc::clone(&finishing); 488 | 489 | executor 490 | .spawn(async move { 491 | started.set(true); 492 | 493 | fut1.await; 494 | 495 | fut1_done.set(true); 496 | 497 | fut2.await; 498 | 499 | finishing.set(true); 500 | }) 501 | .unwrap(); 502 | } 503 | 504 | // not started yet, no progress 505 | assert_eq!(executor.have_tasks(), true); 506 | assert_eq!(started.get(), false); 507 | 508 | executor.run_until_stalled(); 509 | 510 | // started, but fut1 not ready 511 | assert_eq!(executor.have_tasks(), true); 512 | assert_eq!(started.get(), true); 513 | assert_eq!(fut1_done.get(), false); 514 | 515 | handle1.set_ready(); 516 | executor.run_until_stalled(); 517 | 518 | // fut1 finished 519 | assert_eq!(executor.have_tasks(), true); 520 | assert_eq!(fut1_done.get(), true); 521 | assert_eq!(finishing.get(), false); 522 | 523 | handle2.set_ready(); 524 | executor.run_until_stalled(); 525 | 526 | // fut2 finished, and thus the task finished 527 | assert_eq!(finishing.get(), true); 528 | assert_eq!(executor.have_tasks(), false); 529 | } 530 | 531 | #[test] 532 | fn test_executor_run() { 533 | let executor = Executor::new(1); 534 | 535 | let fut = TestFuture::new(); 536 | let handle = fut.handle(); 537 | 538 | executor 539 | .spawn(async move { 540 | fut.await; 541 | }) 542 | .unwrap(); 543 | 544 | executor 545 | .run(|_| { 546 | handle.set_ready(); 547 | 548 | Ok(()) 549 | }) 550 | .unwrap(); 551 | 552 | assert_eq!(executor.have_tasks(), false); 553 | } 554 | 555 | #[test] 556 | fn test_executor_spawn_error() { 557 | let executor = Executor::new(1); 558 | 559 | assert!(executor.spawn(async {}).is_ok()); 560 | assert!(executor.spawn(async {}).is_err()); 561 | } 562 | 563 | #[test] 564 | fn test_executor_current() { 565 | assert!(Executor::current().is_none()); 566 | 567 | let executor = Executor::new(2); 568 | 569 | let flag = Rc::new(Cell::new(false)); 570 | 571 | { 572 | let flag = flag.clone(); 573 | 574 | executor 575 | .spawn(async move { 576 | Executor::current() 577 | .unwrap() 578 | .spawn(async move { 579 | flag.set(true); 580 | }) 581 | .unwrap(); 582 | }) 583 | .unwrap(); 584 | } 585 | 586 | assert_eq!(flag.get(), false); 587 | 588 | executor.run(|_| Ok(())).unwrap(); 589 | 590 | assert_eq!(flag.get(), true); 591 | 592 | let current = Executor::current().unwrap(); 593 | 594 | assert_eq!(executor.have_tasks(), false); 595 | assert!(current.spawn(async {}).is_ok()); 596 | assert_eq!(executor.have_tasks(), true); 597 | 598 | mem::drop(executor); 599 | 600 | assert!(Executor::current().is_some()); 601 | 602 | mem::drop(current); 603 | 604 | assert!(Executor::current().is_none()); 605 | } 606 | 607 | #[test] 608 | fn test_executor_spawner() { 609 | let executor = Executor::new(2); 610 | 611 | let flag = Rc::new(Cell::new(false)); 612 | 613 | { 614 | let flag = flag.clone(); 615 | let spawner = executor.spawner(); 616 | 617 | executor 618 | .spawn(async move { 619 | spawner 620 | .spawn(async move { 621 | flag.set(true); 622 | }) 623 | .unwrap(); 624 | }) 625 | .unwrap(); 626 | } 627 | 628 | assert_eq!(flag.get(), false); 629 | 630 | executor.run(|_| Ok(())).unwrap(); 631 | 632 | assert_eq!(flag.get(), true); 633 | } 634 | 635 | #[test] 636 | fn test_executor_early_wake() { 637 | let executor = Executor::new(1); 638 | 639 | let fut = EarlyWakeFuture::new(); 640 | 641 | executor 642 | .spawn(async move { 643 | fut.await; 644 | }) 645 | .unwrap(); 646 | 647 | let mut park_count = 0; 648 | 649 | executor 650 | .run(|_| { 651 | park_count += 1; 652 | 653 | Ok(()) 654 | }) 655 | .unwrap(); 656 | 657 | assert_eq!(park_count, 1); 658 | } 659 | 660 | #[test] 661 | fn test_executor_pre_poll() { 662 | let executor = Executor::new(1); 663 | 664 | let flag = Rc::new(Cell::new(false)); 665 | 666 | { 667 | let flag = flag.clone(); 668 | 669 | executor.set_pre_poll(move || { 670 | flag.set(true); 671 | }); 672 | } 673 | 674 | executor.spawn(async {}).unwrap(); 675 | 676 | assert_eq!(flag.get(), false); 677 | 678 | executor.run(|_| Ok(())).unwrap(); 679 | 680 | assert_eq!(flag.get(), true); 681 | } 682 | } 683 | -------------------------------------------------------------------------------- /src/timer.rs: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2020-2021 Fanout, Inc. 3 | * 4 | * Licensed under the Apache License, Version 2.0 (the "License"); 5 | * you may not use this file except in compliance with the License. 6 | * You may obtain a copy of the License at 7 | * 8 | * http://www.apache.org/licenses/LICENSE-2.0 9 | * 10 | * Unless required by applicable law or agreed to in writing, software 11 | * distributed under the License is distributed on an "AS IS" BASIS, 12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | * See the License for the specific language governing permissions and 14 | * limitations under the License. 15 | */ 16 | 17 | // adapted from http://25thandclement.com/~william/projects/timeout.c.html (MIT licensed) 18 | 19 | use crate::list; 20 | use slab::Slab; 21 | use std::cmp; 22 | 23 | const WHEEL_BITS: usize = 6; 24 | const WHEEL_NUM: usize = 4; 25 | const WHEEL_LEN: usize = 1 << WHEEL_BITS; 26 | const WHEEL_MAX: usize = WHEEL_LEN - 1; 27 | const WHEEL_MASK: u64 = (WHEEL_LEN as u64) - 1; 28 | const TIMEOUT_MAX: u64 = (1 << (WHEEL_BITS * WHEEL_NUM)) - 1; 29 | 30 | // find last set 31 | fn fls64(x: u64) -> u32 { 32 | 64 - x.leading_zeros() 33 | } 34 | 35 | fn need_resched(curtime: u64, newtime: u64) -> [u64; WHEEL_NUM] { 36 | let mut result = [0; WHEEL_NUM]; 37 | 38 | // no time elapsed 39 | if newtime <= curtime { 40 | return result; 41 | } 42 | 43 | let mut elapsed = newtime - curtime; 44 | 45 | for (wheel, item) in result.iter_mut().enumerate() { 46 | // we only care about the highest bits 47 | let trunc_bits = (wheel * WHEEL_BITS) as u64; 48 | 49 | let pending = if (elapsed >> trunc_bits) > (WHEEL_MAX as u64) { 50 | // all slots need processing 51 | !0 52 | } else { 53 | let old_slot = (curtime >> trunc_bits) & WHEEL_MASK; 54 | let new_slot = (newtime >> trunc_bits) & WHEEL_MASK; 55 | 56 | let d = if new_slot > old_slot { 57 | new_slot - old_slot 58 | } else { 59 | (WHEEL_LEN as u64) - old_slot + new_slot 60 | }; 61 | 62 | if d >= WHEEL_LEN as u64 { 63 | !0 64 | } else if wheel > 0 { 65 | ((1 << d) - 1u64).rotate_left(old_slot as u32) 66 | } else { 67 | ((1 << d) - 1u64).rotate_left((old_slot + 1) as u32) 68 | } 69 | }; 70 | 71 | *item = pending; 72 | 73 | let finished_bit = if wheel > 0 { 74 | // higher wheels have completed a full rotation when slot 63 is processed 75 | 1 << (WHEEL_LEN - 1) 76 | } else { 77 | // lowest wheel has completed a full rotation when slot 0 is processed 78 | 1 79 | }; 80 | 81 | // if the current wheel didn't finish a full rotation then we don't need to look 82 | // at higher wheels 83 | if pending & finished_bit == 0 { 84 | break; 85 | } 86 | 87 | // ensure the elapsed time includes the current slot of the next wheel 88 | elapsed = cmp::max(elapsed, (WHEEL_LEN << (wheel * WHEEL_BITS)) as u64); 89 | } 90 | 91 | result 92 | } 93 | 94 | #[cfg(test)] 95 | fn need_resched_simple(curtime: u64, newtime: u64) -> [u64; WHEEL_NUM] { 96 | let mut result = [0; WHEEL_NUM]; 97 | 98 | // no time elapsed 99 | if newtime <= curtime { 100 | return result; 101 | } 102 | 103 | for curtime in curtime..newtime { 104 | for wheel in 0..WHEEL_NUM { 105 | let trunc_bits = (wheel * WHEEL_BITS) as u64; 106 | 107 | let old_slot = (curtime >> trunc_bits) & WHEEL_MASK; 108 | let new_slot = ((curtime + 1) >> trunc_bits) & WHEEL_MASK; 109 | 110 | if old_slot != new_slot { 111 | if wheel > 0 { 112 | result[wheel] |= 1 << old_slot; 113 | } else { 114 | result[wheel] |= 1 << new_slot; 115 | } 116 | } 117 | } 118 | } 119 | 120 | result 121 | } 122 | 123 | enum InList { 124 | Wheel(usize, usize), 125 | Expired, 126 | } 127 | 128 | struct Timer { 129 | expires: u64, 130 | list: Option, 131 | user_data: usize, 132 | } 133 | 134 | pub struct TimerWheel { 135 | nodes: Slab>, 136 | wheel: [[list::List; WHEEL_LEN]; WHEEL_NUM], 137 | expired: list::List, 138 | pending: [u64; WHEEL_NUM], 139 | curtime: u64, 140 | } 141 | 142 | impl TimerWheel { 143 | pub fn new(capacity: usize) -> Self { 144 | Self { 145 | nodes: Slab::with_capacity(capacity), 146 | wheel: [[list::List::default(); WHEEL_LEN]; WHEEL_NUM], 147 | expired: list::List::default(), 148 | pending: [0; WHEEL_NUM], 149 | curtime: 0, 150 | } 151 | } 152 | 153 | #[allow(clippy::result_unit_err)] 154 | pub fn add(&mut self, expires: u64, user_data: usize) -> Result { 155 | if self.nodes.len() == self.nodes.capacity() { 156 | return Err(()); 157 | } 158 | 159 | let t = Timer { 160 | expires, 161 | list: None, 162 | user_data, 163 | }; 164 | 165 | let key = self.nodes.insert(list::Node::new(t)); 166 | 167 | self.sched(key); 168 | 169 | Ok(key) 170 | } 171 | 172 | pub fn remove(&mut self, key: usize) { 173 | let n = match self.nodes.get(key) { 174 | Some(n) => n, 175 | None => return, 176 | }; 177 | 178 | match n.value.list { 179 | Some(InList::Wheel(wheel, slot)) => { 180 | let l = &mut self.wheel[wheel][slot]; 181 | l.remove(&mut self.nodes, key); 182 | 183 | if l.is_empty() { 184 | self.pending[wheel] &= !(1 << slot); 185 | } 186 | } 187 | Some(InList::Expired) => { 188 | self.expired.remove(&mut self.nodes, key); 189 | } 190 | None => {} 191 | } 192 | 193 | self.nodes.remove(key); 194 | } 195 | 196 | pub fn timeout(&self) -> Option { 197 | if !self.expired.is_empty() { 198 | return Some(0); 199 | } 200 | 201 | let mut timeout = None; 202 | let mut relmask = 0; 203 | 204 | for wheel in 0..WHEEL_NUM { 205 | // we only care about the highest bits 206 | let trunc_bits = (wheel * WHEEL_BITS) as u64; 207 | 208 | if self.pending[wheel] != 0 { 209 | let slot = ((self.curtime >> trunc_bits) & WHEEL_MASK) as usize; 210 | 211 | let pending = self.pending[wheel].rotate_right(slot as u32); 212 | 213 | // for higher order wheels, timeouts are one step in the future 214 | let offset = if wheel > 0 { 1 } else { 0 }; 215 | 216 | // pending is guaranteed to be non-zero 217 | let t = ((pending.trailing_zeros() as u64) + offset) << trunc_bits; 218 | 219 | // reduce by how much lower wheels have progressed 220 | let t = t - (relmask & self.curtime); 221 | 222 | timeout = Some(match timeout { 223 | Some(best) => cmp::min(best, t), 224 | None => t, 225 | }); 226 | } 227 | 228 | relmask <<= WHEEL_BITS; 229 | relmask |= WHEEL_MASK; 230 | } 231 | 232 | timeout 233 | } 234 | 235 | pub fn update(&mut self, curtime: u64) { 236 | // time must go forward 237 | if curtime <= self.curtime { 238 | return; 239 | } 240 | 241 | let need = need_resched(self.curtime, curtime); 242 | 243 | let mut l = list::List::default(); 244 | 245 | for (wheel, &pending) in need.iter().enumerate() { 246 | // loop as long as we still have slots to process 247 | while pending & self.pending[wheel] != 0 { 248 | // get rightmost (earliest) slot that needs processing 249 | let slot = (pending & self.pending[wheel]).trailing_zeros() as usize; 250 | 251 | // move the timers out 252 | l.concat(&mut self.nodes, &mut self.wheel[wheel][slot]); 253 | self.pending[wheel] &= !(1 << slot); 254 | } 255 | } 256 | 257 | self.curtime = curtime; 258 | 259 | while let Some(key) = l.head { 260 | l.remove(&mut self.nodes, key); 261 | 262 | let n = &mut self.nodes[key]; 263 | n.value.list = None; 264 | 265 | self.sched(key); 266 | } 267 | } 268 | 269 | pub fn take_expired(&mut self) -> Option<(usize, usize)> { 270 | match self.expired.pop_front(&mut self.nodes) { 271 | Some(key) => { 272 | let n = &self.nodes[key]; 273 | let user_data = n.value.user_data; 274 | 275 | self.nodes.remove(key); 276 | 277 | Some((key, user_data)) 278 | } 279 | None => None, 280 | } 281 | } 282 | 283 | fn sched(&mut self, key: usize) { 284 | let n = &self.nodes[key]; 285 | let expires = n.value.expires; 286 | 287 | if expires > self.curtime { 288 | // get relative timeout, capped 289 | let t = cmp::min(expires - self.curtime, TIMEOUT_MAX); 290 | assert!(t > 0); 291 | 292 | // wheel is selected by relative time 293 | // t = 0 = not valid 294 | // t = 1 = 0b0_000000_000001 -> fls 1, wheel 0 295 | // t = 63 = 0b0_000000_111111 -> fls 6, wheel 0 296 | // t = 64 = 0b0_000001_000000 -> fls 7, wheel 1 297 | // t = 4032 = 0b0_111111_000000 -> fls 12, wheel 1 298 | // t = 4095 = 0b0_111111_111111 -> fls 12, wheel 1 299 | // t = 4096 = 0b1_000000_000000 -> fls 13, wheel 2 300 | let wheel = ((fls64(t) - 1) as usize) / WHEEL_BITS; 301 | assert!(wheel < WHEEL_NUM); 302 | 303 | // we only care about the highest bits 304 | let trunc_bits = (wheel * WHEEL_BITS) as u64; 305 | 306 | // for higher order wheels, schedule 1 slot early. this way, fractional 307 | // time remaining can be rescheduled to a lower wheel 308 | let offset = if wheel > 0 { 1 } else { 0 }; 309 | 310 | // slot is selected by absolute time 311 | let slot = (((expires >> trunc_bits) - offset) & WHEEL_MASK) as usize; 312 | 313 | self.wheel[wheel][slot].push_back(&mut self.nodes, key); 314 | self.pending[wheel] |= 1 << slot; 315 | 316 | let n = &mut self.nodes[key]; 317 | n.value.list = Some(InList::Wheel(wheel, slot)); 318 | } else { 319 | self.expired.push_back(&mut self.nodes, key); 320 | 321 | let n = &mut self.nodes[key]; 322 | n.value.list = Some(InList::Expired); 323 | } 324 | } 325 | } 326 | 327 | #[cfg(test)] 328 | mod tests { 329 | use super::*; 330 | 331 | // convert string time of the form "x:x:x:x", where each part is a number between 0-63 332 | fn ts(s: &str) -> u64 { 333 | let mut result = 0; 334 | 335 | for (i, part) in s.rsplit(":").enumerate() { 336 | let x: u64 = part.parse().unwrap(); 337 | assert!(x <= (WHEEL_MAX as u64)); 338 | 339 | result |= x << (i * WHEEL_BITS); 340 | } 341 | 342 | result 343 | } 344 | 345 | // convert string range to bits 346 | fn r2b(s: &str) -> u64 { 347 | let mut it = s.split("-"); 348 | 349 | let start = it.next().unwrap(); 350 | let end = it.next().unwrap(); 351 | 352 | assert_eq!(it.next(), None); 353 | 354 | let mut pos: u64 = start.parse().unwrap(); 355 | let end: u64 = end.parse().unwrap(); 356 | 357 | let mut result = 0; 358 | 359 | loop { 360 | result |= 1 << pos; 361 | 362 | if pos == end { 363 | break; 364 | } 365 | 366 | pos = (pos + 1) & WHEEL_MASK; 367 | } 368 | 369 | result 370 | } 371 | 372 | // convert wheel ranges of the form "x:x:x:x", where each part is a range 373 | fn r2w(s: &str) -> [u64; WHEEL_NUM] { 374 | let mut result = [0; WHEEL_NUM]; 375 | 376 | for (i, part) in s.rsplit(":").enumerate() { 377 | if !part.is_empty() { 378 | result[i] = r2b(part); 379 | } 380 | } 381 | 382 | result 383 | } 384 | 385 | #[test] 386 | fn test_fls() { 387 | assert_eq!(fls64(0), 0); 388 | assert_eq!(fls64(0b1), 1); 389 | assert_eq!(fls64(0b10), 2); 390 | assert_eq!(fls64(0x4000000000000000), 63); 391 | assert_eq!(fls64(0x8000000000000000), 64); 392 | } 393 | 394 | #[test] 395 | fn test_sched() { 396 | let mut w = TimerWheel::new(10); 397 | 398 | w.update(7); 399 | 400 | // expired 401 | let t1 = w.add(0b0_000000, 1).unwrap(); 402 | 403 | // wheel 0 slot 8 (1 tick away) 404 | let t2 = w.add(0b0_001000, 1).unwrap(); 405 | 406 | // wheel 0 slot 63 (56 ticks away) 407 | let t3 = w.add(0b0_111111, 1).unwrap(); 408 | 409 | // wheel 0 slot 0 (57 ticks away) 410 | let t4 = w.add(0b1_000000, 1).unwrap(); 411 | 412 | // wheel 1 slot 0 413 | let t5 = w.add(0b1_001000, 1).unwrap(); 414 | 415 | // wheel 3 slot 63 416 | let t6 = w.add(0b1_000000_000000_000000_000000, 1).unwrap(); 417 | 418 | assert_eq!(w.expired.head, Some(t1)); 419 | assert_eq!(w.wheel[0][8].head, Some(t2)); 420 | assert_eq!(w.wheel[0][63].head, Some(t3)); 421 | assert_eq!(w.wheel[0][0].head, Some(t4)); 422 | assert_eq!(w.wheel[1][0].head, Some(t5)); 423 | assert_eq!(w.wheel[3][63].head, Some(t6)); 424 | } 425 | 426 | #[test] 427 | fn test_need_resched() { 428 | struct Test { 429 | curtime: &'static str, 430 | newtime: &'static str, 431 | expected: &'static str, 432 | } 433 | 434 | fn t(curtime: &'static str, newtime: &'static str, expected: &'static str) -> Test { 435 | Test { 436 | curtime, 437 | newtime, 438 | expected, 439 | } 440 | } 441 | 442 | let table = [ 443 | t("00:00", "00:00", ""), 444 | t("00:00", "00:01", "01-01"), 445 | t("00:01", "00:02", "02-02"), 446 | t("00:02", "00:63", "03-63"), 447 | t("00:63", "01:00", "00-00:00-00"), 448 | t("01:00", "01:02", "01-02"), 449 | t("01:02", "05:01", "01-04:00-63"), 450 | t("05:01", "05:02", "02-02"), 451 | t("05:02", "06:01", "05-05:03-01"), 452 | t("00:63:63", "01:00:00", "00-00:63-63:00-00"), 453 | t("08:00:00", "08:01:00", "00-00:00-63"), 454 | t("04:00:02", "05:01:00", "04-04:00-63:00-63"), 455 | t("04:01:02", "05:00:00", "04-04:01-63:00-63"), 456 | t("04:00:03", "05:00:00", "04-04:00-63:00-63"), 457 | t("04:00:02", "05:00:00", "04-04:00-63:00-63"), 458 | t("08:00:19", "08:62:63", "00-61:00-63"), 459 | t("08:00:19", "08:63:63", "00-62:00-63"), 460 | t("09:00:00", "09:63:62", "00-62:00-63"), 461 | ]; 462 | 463 | for (row, t) in table.iter().enumerate() { 464 | let curtime = ts(t.curtime); 465 | let newtime = ts(t.newtime); 466 | let expected = r2w(t.expected); 467 | 468 | // ensure the simple algorithm returns what we expect 469 | assert_eq!( 470 | need_resched_simple(curtime, newtime), 471 | expected, 472 | "row={} curtime={} newtime={}", 473 | row, 474 | curtime, 475 | newtime 476 | ); 477 | 478 | // ensure the optimized algorithm returns matching results 479 | assert_eq!( 480 | need_resched(curtime, newtime), 481 | expected, 482 | "row={} curtime={} newtime={}", 483 | row, 484 | curtime, 485 | newtime 486 | ); 487 | } 488 | } 489 | 490 | #[test] 491 | fn test_rotate() { 492 | // test full rotations through wheels 0 and 1, and one step of wheel 2 493 | let count = (64 * 64) + 1; 494 | 495 | let mut w = TimerWheel::new(count); 496 | 497 | for i in 0..count { 498 | w.add(i as u64, i).unwrap(); 499 | } 500 | 501 | for i in 0..count { 502 | let (_, v) = w.take_expired().unwrap(); 503 | assert_eq!(v, i); 504 | assert_eq!(w.take_expired(), None); 505 | 506 | w.update((i + 1) as u64); 507 | } 508 | 509 | assert_eq!(w.take_expired(), None); 510 | } 511 | 512 | #[test] 513 | fn test_wheel() { 514 | let mut w = TimerWheel::new(10); 515 | 516 | assert_eq!(w.timeout(), None); 517 | assert_eq!(w.take_expired(), None); 518 | 519 | let t1 = w.add(4, 1).unwrap(); 520 | assert_eq!(w.timeout(), Some(4)); 521 | 522 | w.remove(t1); 523 | assert_eq!(w.timeout(), None); 524 | 525 | w.update(5); 526 | assert_eq!(w.take_expired(), None); 527 | 528 | let t2 = w.add(8, 2).unwrap(); 529 | assert_eq!(w.timeout(), Some(3)); 530 | 531 | w.update(7); 532 | assert_eq!(w.timeout(), Some(1)); 533 | assert_eq!(w.take_expired(), None); 534 | 535 | w.update(8); 536 | assert_eq!(w.timeout(), Some(0)); 537 | assert_eq!(w.take_expired(), Some((t2, 2))); 538 | assert_eq!(w.take_expired(), None); 539 | 540 | for i in 0..2 { 541 | let base = i * 20_000_000; 542 | 543 | let t1 = w.add(base + 1, 1).unwrap(); 544 | let t2 = w.add(base + 10, 2).unwrap(); 545 | let t3 = w.add(base + 1_000, 3).unwrap(); 546 | let t4 = w.add(base + 100_000, 4).unwrap(); 547 | let t5 = w.add(base + 10_000_000, 5).unwrap(); 548 | 549 | w.update(base + 100); 550 | assert_eq!(w.timeout(), Some(0)); 551 | assert_eq!(w.take_expired(), Some((t1, 1))); 552 | assert_eq!(w.take_expired(), Some((t2, 2))); 553 | assert_eq!(w.take_expired(), None); 554 | assert!(w.timeout().unwrap() <= 900); 555 | 556 | w.update(base + 2_000); 557 | assert_eq!(w.timeout(), Some(0)); 558 | assert_eq!(w.take_expired(), Some((t3, 3))); 559 | assert_eq!(w.take_expired(), None); 560 | assert!(w.timeout().unwrap() <= 98_000); 561 | 562 | w.update(base + 200_000); 563 | assert_eq!(w.timeout(), Some(0)); 564 | assert_eq!(w.take_expired(), Some((t4, 4))); 565 | assert_eq!(w.take_expired(), None); 566 | assert!(w.timeout().unwrap() <= 9_800_000); 567 | 568 | w.update(base + 12_000_000); 569 | assert_eq!(w.timeout(), Some(0)); 570 | assert_eq!(w.take_expired(), Some((t5, 5))); 571 | assert_eq!(w.take_expired(), None); 572 | assert_eq!(w.timeout(), None); 573 | } 574 | } 575 | } 576 | -------------------------------------------------------------------------------- /src/arena.rs: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2020-2023 Fanout, Inc. 3 | * 4 | * Licensed under the Apache License, Version 2.0 (the "License"); 5 | * you may not use this file except in compliance with the License. 6 | * You may obtain a copy of the License at 7 | * 8 | * http://www.apache.org/licenses/LICENSE-2.0 9 | * 10 | * Unless required by applicable law or agreed to in writing, software 11 | * distributed under the License is distributed on an "AS IS" BASIS, 12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | * See the License for the specific language governing permissions and 14 | * limitations under the License. 15 | */ 16 | 17 | use slab::Slab; 18 | use std::cell::{RefCell, RefMut}; 19 | use std::mem; 20 | use std::ops::{Deref, DerefMut}; 21 | use std::sync::{Mutex, MutexGuard}; 22 | 23 | pub struct EntryGuard<'a, T> { 24 | entries: RefMut<'a, Slab>, 25 | entry: &'a mut T, 26 | key: usize, 27 | } 28 | 29 | impl EntryGuard<'_, T> { 30 | fn remove(mut self) { 31 | self.entries.remove(self.key); 32 | } 33 | } 34 | 35 | impl Deref for EntryGuard<'_, T> { 36 | type Target = T; 37 | 38 | fn deref(&self) -> &Self::Target { 39 | self.entry 40 | } 41 | } 42 | 43 | impl DerefMut for EntryGuard<'_, T> { 44 | fn deref_mut(&mut self) -> &mut Self::Target { 45 | self.entry 46 | } 47 | } 48 | 49 | // this is essentially a sharable slab for use within a single thread. 50 | // operations are protected by a RefCell. when an element is retrieved for 51 | // reading or modification, it is wrapped in a EntryGuard which keeps the 52 | // entire slab borrowed until the caller is done working with the element 53 | pub struct Memory { 54 | entries: RefCell>, 55 | } 56 | 57 | impl Memory { 58 | pub fn new(capacity: usize) -> Self { 59 | // allocate the slab with fixed capacity 60 | let s = Slab::with_capacity(capacity); 61 | 62 | Self { 63 | entries: RefCell::new(s), 64 | } 65 | } 66 | 67 | #[cfg(test)] 68 | pub fn len(&self) -> usize { 69 | let entries = self.entries.borrow(); 70 | 71 | entries.len() 72 | } 73 | 74 | fn insert(&self, e: T) -> Result { 75 | let mut entries = self.entries.borrow_mut(); 76 | 77 | // out of capacity. by preventing inserts beyond the capacity, we 78 | // ensure the underlying memory won't get moved due to a realloc 79 | if entries.len() == entries.capacity() { 80 | return Err(()); 81 | } 82 | 83 | Ok(entries.insert(e)) 84 | } 85 | 86 | fn get<'a>(&'a self, key: usize) -> Option> { 87 | let mut entries = self.entries.borrow_mut(); 88 | 89 | let entry = entries.get_mut(key)?; 90 | 91 | // slab element addresses are guaranteed to be stable once created, 92 | // and the only place we remove the element is in EntryGuard's 93 | // remove method which consumes itself, therefore it is safe to 94 | // assume the element will live at least as long as the EntryGuard 95 | // and we can extend the lifetime of the reference beyond the 96 | // RefMut 97 | let entry = unsafe { mem::transmute::<&mut T, &'a mut T>(entry) }; 98 | 99 | Some(EntryGuard { 100 | entries, 101 | entry, 102 | key, 103 | }) 104 | } 105 | 106 | // for tests, as a way to confirm the memory isn't moving. be careful 107 | // with this. the very first element inserted will be at index 0, but 108 | // if the slab has been used and cleared, then the next element 109 | // inserted may not be at index 0 and calling this method afterward 110 | // will panic 111 | #[cfg(test)] 112 | fn entry0_ptr(&self) -> *const T { 113 | let entries = self.entries.borrow(); 114 | 115 | entries.get(0).unwrap() as *const T 116 | } 117 | } 118 | 119 | pub struct SyncEntryGuard<'a, T> { 120 | entries: MutexGuard<'a, Slab>, 121 | entry: &'a mut T, 122 | key: usize, 123 | } 124 | 125 | impl SyncEntryGuard<'_, T> { 126 | fn remove(mut self) { 127 | self.entries.remove(self.key); 128 | } 129 | } 130 | 131 | impl Deref for SyncEntryGuard<'_, T> { 132 | type Target = T; 133 | 134 | fn deref(&self) -> &Self::Target { 135 | self.entry 136 | } 137 | } 138 | 139 | impl DerefMut for SyncEntryGuard<'_, T> { 140 | fn deref_mut(&mut self) -> &mut Self::Target { 141 | self.entry 142 | } 143 | } 144 | 145 | // this is essentially a thread-safe slab. operations are protected by a 146 | // mutex. when an element is retrieved for reading or modification, it is 147 | // wrapped in a EntryGuard which keeps the entire slab locked until the 148 | // caller is done working with the element 149 | pub struct SyncMemory { 150 | entries: Mutex>, 151 | } 152 | 153 | impl SyncMemory { 154 | pub fn new(capacity: usize) -> Self { 155 | // allocate the slab with fixed capacity 156 | let s = Slab::with_capacity(capacity); 157 | 158 | Self { 159 | entries: Mutex::new(s), 160 | } 161 | } 162 | 163 | #[cfg(test)] 164 | pub fn len(&self) -> usize { 165 | let entries = self.entries.lock().unwrap(); 166 | 167 | entries.len() 168 | } 169 | 170 | fn insert(&self, e: T) -> Result { 171 | let mut entries = self.entries.lock().unwrap(); 172 | 173 | // out of capacity. by preventing inserts beyond the capacity, we 174 | // ensure the underlying memory won't get moved due to a realloc 175 | if entries.len() == entries.capacity() { 176 | return Err(()); 177 | } 178 | 179 | Ok(entries.insert(e)) 180 | } 181 | 182 | fn get<'a>(&'a self, key: usize) -> Option> { 183 | let mut entries = self.entries.lock().unwrap(); 184 | 185 | let entry = entries.get_mut(key)?; 186 | 187 | // slab element addresses are guaranteed to be stable once created, 188 | // and the only place we remove the element is in SyncEntryGuard's 189 | // remove method which consumes itself, therefore it is safe to 190 | // assume the element will live at least as long as the SyncEntryGuard 191 | // and we can extend the lifetime of the reference beyond the 192 | // MutexGuard 193 | let entry = unsafe { mem::transmute::<&mut T, &'a mut T>(entry) }; 194 | 195 | Some(SyncEntryGuard { 196 | entries, 197 | entry, 198 | key, 199 | }) 200 | } 201 | 202 | // for tests, as a way to confirm the memory isn't moving. be careful 203 | // with this. the very first element inserted will be at index 0, but 204 | // if the slab has been used and cleared, then the next element 205 | // inserted may not be at index 0 and calling this method afterward 206 | // will panic 207 | #[cfg(test)] 208 | fn entry0_ptr(&self) -> *const T { 209 | let entries = self.entries.lock().unwrap(); 210 | 211 | entries.get(0).unwrap() as *const T 212 | } 213 | } 214 | 215 | pub struct ReusableValue { 216 | reusable: std::sync::Arc>, 217 | value: *mut T, 218 | key: usize, 219 | } 220 | 221 | impl ReusableValue { 222 | // vec element addresses are guaranteed to be stable once created, 223 | // and elements are only removed when the Reusable is dropped, and 224 | // the Arc'd Reusable is guaranteed to live as long as 225 | // ReusableValue, therefore it is safe to assume the element will 226 | // live at least as long as the ReusableValue 227 | 228 | fn get(&self) -> &T { 229 | unsafe { &*self.value } 230 | } 231 | 232 | fn get_mut(&mut self) -> &mut T { 233 | unsafe { &mut *self.value } 234 | } 235 | } 236 | 237 | impl Drop for ReusableValue { 238 | fn drop(&mut self) { 239 | let mut entries = self.reusable.entries.lock().unwrap(); 240 | 241 | entries.0.remove(self.key); 242 | } 243 | } 244 | 245 | impl Deref for ReusableValue { 246 | type Target = T; 247 | 248 | fn deref(&self) -> &Self::Target { 249 | self.get() 250 | } 251 | } 252 | 253 | impl DerefMut for ReusableValue { 254 | fn deref_mut(&mut self) -> &mut Self::Target { 255 | self.get_mut() 256 | } 257 | } 258 | 259 | // like Memory, but for preinitializing each value and reusing 260 | pub struct Reusable { 261 | entries: Mutex<(Slab<()>, Vec)>, 262 | } 263 | 264 | impl Reusable { 265 | pub fn new(capacity: usize, init_fn: F) -> Self 266 | where 267 | F: Fn() -> T, 268 | { 269 | let mut values = Vec::with_capacity(capacity); 270 | 271 | for _ in 0..capacity { 272 | values.push(init_fn()); 273 | } 274 | 275 | // allocate the slab with fixed capacity 276 | let s = Slab::with_capacity(capacity); 277 | 278 | Self { 279 | entries: Mutex::new((s, values)), 280 | } 281 | } 282 | 283 | #[cfg(test)] 284 | pub fn len(&self) -> usize { 285 | let entries = self.entries.lock().unwrap(); 286 | 287 | entries.0.len() 288 | } 289 | 290 | #[allow(clippy::result_unit_err)] 291 | pub fn reserve(self: &std::sync::Arc) -> Result, ()> { 292 | let mut entries = self.entries.lock().unwrap(); 293 | 294 | // out of capacity. the number of buffers is fixed 295 | if entries.0.len() == entries.0.capacity() { 296 | return Err(()); 297 | } 298 | 299 | let key = entries.0.insert(()); 300 | 301 | let value = &mut entries.1[key] as *mut T; 302 | 303 | Ok(ReusableValue { 304 | reusable: self.clone(), 305 | value, 306 | key, 307 | }) 308 | } 309 | } 310 | 311 | pub struct RcEntry { 312 | value: T, 313 | refs: usize, 314 | } 315 | 316 | pub type RcMemory = Memory>; 317 | 318 | pub struct Rc { 319 | memory: std::rc::Rc>, 320 | key: usize, 321 | } 322 | 323 | impl Rc { 324 | #[allow(clippy::result_unit_err)] 325 | pub fn new(v: T, memory: &std::rc::Rc>) -> Result { 326 | let key = memory.insert(RcEntry { value: v, refs: 1 })?; 327 | 328 | Ok(Self { 329 | memory: std::rc::Rc::clone(memory), 330 | key, 331 | }) 332 | } 333 | 334 | #[allow(clippy::should_implement_trait)] 335 | pub fn clone(rc: &Rc) -> Self { 336 | let mut e = rc.memory.get(rc.key).unwrap(); 337 | 338 | e.refs += 1; 339 | 340 | Self { 341 | memory: rc.memory.clone(), 342 | key: rc.key, 343 | } 344 | } 345 | 346 | pub fn get<'a>(&'a self) -> &'a T { 347 | let e = self.memory.get(self.key).unwrap(); 348 | 349 | // get a reference to the inner value 350 | let value = &e.value; 351 | 352 | // entry addresses are guaranteed to be stable once created, and the 353 | // entry managed by this Rc won't be dropped until this Rc drops, 354 | // therefore it is safe to assume the entry managed by this Rc will 355 | // live at least as long as this Rc, and we can extend the lifetime 356 | // of the reference beyond the EntryGuard 357 | unsafe { mem::transmute::<&T, &'a T>(value) } 358 | } 359 | } 360 | 361 | impl Drop for Rc { 362 | fn drop(&mut self) { 363 | let mut e = self.memory.get(self.key).unwrap(); 364 | 365 | if e.refs == 1 { 366 | e.remove(); 367 | return; 368 | } 369 | 370 | e.refs -= 1; 371 | } 372 | } 373 | 374 | pub type ArcMemory = SyncMemory>; 375 | 376 | pub struct Arc { 377 | memory: std::sync::Arc>, 378 | key: usize, 379 | } 380 | 381 | impl Arc { 382 | #[allow(clippy::result_unit_err)] 383 | pub fn new(v: T, memory: &std::sync::Arc>) -> Result { 384 | let key = memory.insert(RcEntry { value: v, refs: 1 })?; 385 | 386 | Ok(Self { 387 | memory: memory.clone(), 388 | key, 389 | }) 390 | } 391 | 392 | #[allow(clippy::should_implement_trait)] 393 | pub fn clone(rc: &Arc) -> Self { 394 | let mut e = rc.memory.get(rc.key).unwrap(); 395 | 396 | e.refs += 1; 397 | 398 | Self { 399 | memory: rc.memory.clone(), 400 | key: rc.key, 401 | } 402 | } 403 | 404 | pub fn get<'a>(&'a self) -> &'a T { 405 | let e = self.memory.get(self.key).unwrap(); 406 | 407 | // get a reference to the inner value 408 | let value = &e.value; 409 | 410 | // entry addresses are guaranteed to be stable once created, and the 411 | // entry managed by this Arc won't be dropped until this Arc drops, 412 | // therefore it is safe to assume the entry managed by this Arc will 413 | // live at least as long as this Arc, and we can extend the lifetime 414 | // of the reference beyond the SyncEntryGuard 415 | unsafe { mem::transmute::<&T, &'a T>(value) } 416 | } 417 | } 418 | 419 | impl Drop for Arc { 420 | fn drop(&mut self) { 421 | let mut e = self.memory.get(self.key).unwrap(); 422 | 423 | if e.refs == 1 { 424 | e.remove(); 425 | return; 426 | } 427 | 428 | e.refs -= 1; 429 | } 430 | } 431 | 432 | // adapted from https://github.com/rust-lang/rfcs/pull/2802 433 | pub fn recycle_vec(mut v: Vec) -> Vec { 434 | assert_eq!(core::mem::size_of::(), core::mem::size_of::()); 435 | assert_eq!(core::mem::align_of::(), core::mem::align_of::()); 436 | v.clear(); 437 | let ptr = v.as_mut_ptr(); 438 | let capacity = v.capacity(); 439 | mem::forget(v); 440 | let ptr = ptr as *mut U; 441 | unsafe { Vec::from_raw_parts(ptr, 0, capacity) } 442 | } 443 | 444 | // ReusableVec inspired by recycle_vec 445 | 446 | pub struct ReusableVecHandle<'a, T> { 447 | vec: &'a mut Vec, 448 | } 449 | 450 | impl ReusableVecHandle<'_, T> { 451 | pub fn get_ref(&self) -> &Vec { 452 | self.vec 453 | } 454 | 455 | pub fn get_mut(&mut self) -> &mut Vec { 456 | self.vec 457 | } 458 | } 459 | 460 | impl Drop for ReusableVecHandle<'_, T> { 461 | fn drop(&mut self) { 462 | self.vec.clear(); 463 | } 464 | } 465 | 466 | impl Deref for ReusableVecHandle<'_, T> { 467 | type Target = Vec; 468 | 469 | fn deref(&self) -> &Self::Target { 470 | self.get_ref() 471 | } 472 | } 473 | 474 | impl DerefMut for ReusableVecHandle<'_, T> { 475 | fn deref_mut(&mut self) -> &mut Self::Target { 476 | self.get_mut() 477 | } 478 | } 479 | 480 | pub struct ReusableVec { 481 | vec: Vec<()>, 482 | size: usize, 483 | align: usize, 484 | } 485 | 486 | impl ReusableVec { 487 | pub fn new(capacity: usize) -> Self { 488 | let size = mem::size_of::(); 489 | let align = mem::align_of::(); 490 | 491 | let vec: Vec = Vec::with_capacity(capacity); 492 | 493 | // safety: we must cast to Vec before using, where U has the same 494 | // size and alignment as T 495 | let vec: Vec<()> = unsafe { mem::transmute(vec) }; 496 | 497 | Self { vec, size, align } 498 | } 499 | 500 | pub fn get_as_new(&mut self) -> ReusableVecHandle<'_, U> { 501 | let size = mem::size_of::(); 502 | let align = mem::align_of::(); 503 | 504 | // if these don't match, panic. it's up the user to ensure the type 505 | // is acceptable 506 | assert_eq!(self.size, size); 507 | assert_eq!(self.align, align); 508 | 509 | let vec: &mut Vec<()> = &mut self.vec; 510 | 511 | // safety: U has the expected size and alignment 512 | let vec: &mut Vec = unsafe { mem::transmute(vec) }; 513 | 514 | // the vec starts empty, and is always cleared when the handle drops. 515 | // get_as_new() borrows self mutably, so it's not possible to create 516 | // a handle when one already exists 517 | assert!(vec.is_empty()); 518 | 519 | ReusableVecHandle { vec } 520 | } 521 | } 522 | 523 | #[cfg(test)] 524 | mod tests { 525 | use super::*; 526 | 527 | #[test] 528 | fn test_reusable() { 529 | let reusable = std::sync::Arc::new(Reusable::new(2, || vec![0; 128])); 530 | assert_eq!(reusable.len(), 0); 531 | 532 | let mut buf1 = reusable.reserve().unwrap(); 533 | assert_eq!(reusable.len(), 1); 534 | 535 | let mut buf2 = reusable.reserve().unwrap(); 536 | assert_eq!(reusable.len(), 2); 537 | 538 | // no room 539 | assert!(reusable.reserve().is_err()); 540 | 541 | buf1[..5].copy_from_slice(b"hello"); 542 | buf2[..5].copy_from_slice(b"world"); 543 | 544 | assert_eq!(&buf1[..5], b"hello"); 545 | assert_eq!(&buf2[..5], b"world"); 546 | 547 | mem::drop(buf1); 548 | assert_eq!(reusable.len(), 1); 549 | 550 | mem::drop(buf2); 551 | assert_eq!(reusable.len(), 0); 552 | } 553 | 554 | #[test] 555 | fn test_rc() { 556 | let memory = std::rc::Rc::new(RcMemory::new(2)); 557 | assert_eq!(memory.len(), 0); 558 | 559 | let e0a = Rc::new(123 as i32, &memory).unwrap(); 560 | assert_eq!(memory.len(), 1); 561 | let p = memory.entry0_ptr(); 562 | 563 | let e0b = Rc::clone(&e0a); 564 | assert_eq!(memory.len(), 1); 565 | assert_eq!(memory.entry0_ptr(), p); 566 | 567 | let e1a = Rc::new(456 as i32, &memory).unwrap(); 568 | assert_eq!(memory.len(), 2); 569 | assert_eq!(memory.entry0_ptr(), p); 570 | 571 | // no room 572 | assert!(Rc::new(789 as i32, &memory).is_err()); 573 | 574 | assert_eq!(*e0a.get(), 123); 575 | assert_eq!(*e0b.get(), 123); 576 | assert_eq!(*e1a.get(), 456); 577 | 578 | mem::drop(e0b); 579 | assert_eq!(memory.len(), 2); 580 | assert_eq!(memory.entry0_ptr(), p); 581 | 582 | mem::drop(e0a); 583 | assert_eq!(memory.len(), 1); 584 | 585 | mem::drop(e1a); 586 | assert_eq!(memory.len(), 0); 587 | } 588 | 589 | #[test] 590 | fn test_arc() { 591 | let memory = std::sync::Arc::new(ArcMemory::new(2)); 592 | assert_eq!(memory.len(), 0); 593 | 594 | let e0a = Arc::new(123 as i32, &memory).unwrap(); 595 | assert_eq!(memory.len(), 1); 596 | let p = memory.entry0_ptr(); 597 | 598 | let e0b = Arc::clone(&e0a); 599 | assert_eq!(memory.len(), 1); 600 | assert_eq!(memory.entry0_ptr(), p); 601 | 602 | let e1a = Arc::new(456 as i32, &memory).unwrap(); 603 | assert_eq!(memory.len(), 2); 604 | assert_eq!(memory.entry0_ptr(), p); 605 | 606 | // no room 607 | assert!(Arc::new(789 as i32, &memory).is_err()); 608 | 609 | assert_eq!(*e0a.get(), 123); 610 | assert_eq!(*e0b.get(), 123); 611 | assert_eq!(*e1a.get(), 456); 612 | 613 | mem::drop(e0b); 614 | assert_eq!(memory.len(), 2); 615 | assert_eq!(memory.entry0_ptr(), p); 616 | 617 | mem::drop(e0a); 618 | assert_eq!(memory.len(), 1); 619 | 620 | mem::drop(e1a); 621 | assert_eq!(memory.len(), 0); 622 | } 623 | 624 | #[test] 625 | fn test_reusable_vec() { 626 | let mut vec_mem = ReusableVec::new::(100); 627 | 628 | let mut vec = vec_mem.get_as_new::(); 629 | 630 | assert_eq!(vec.capacity(), 100); 631 | assert_eq!(vec.len(), 0); 632 | 633 | vec.push(1); 634 | assert_eq!(vec.len(), 1); 635 | 636 | mem::drop(vec); 637 | 638 | let vec = vec_mem.get_as_new::(); 639 | 640 | assert_eq!(vec.capacity(), 100); 641 | assert_eq!(vec.len(), 0); 642 | } 643 | } 644 | -------------------------------------------------------------------------------- /src/main.rs: -------------------------------------------------------------------------------- 1 | /* 2 | * Copyright (C) 2020-2023 Fanout, Inc. 3 | * 4 | * Licensed under the Apache License, Version 2.0 (the "License"); 5 | * you may not use this file except in compliance with the License. 6 | * You may obtain a copy of the License at 7 | * 8 | * http://www.apache.org/licenses/LICENSE-2.0 9 | * 10 | * Unless required by applicable law or agreed to in writing, software 11 | * distributed under the License is distributed on an "AS IS" BASIS, 12 | * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 | * See the License for the specific language governing permissions and 14 | * limitations under the License. 15 | */ 16 | 17 | use clap::{crate_version, Arg, ArgAction, Command}; 18 | use condure::app; 19 | use log::{error, Level, LevelFilter, Metadata, Record}; 20 | use std::error::Error; 21 | use std::io; 22 | use std::mem; 23 | use std::path::PathBuf; 24 | use std::process; 25 | use std::str; 26 | use std::sync::Once; 27 | use std::time::Duration; 28 | use time::macros::format_description; 29 | use time::{OffsetDateTime, UtcOffset}; 30 | 31 | // safety values 32 | const WORKERS_MAX: usize = 1024; 33 | const CONNS_MAX: usize = 10_000_000; 34 | 35 | const PRIVATE_SUBNETS: &[&str] = &[ 36 | "127.0.0.0/8", 37 | "10.0.0.0/8", 38 | "172.16.0.0/12", 39 | "192.168.0.0/16", 40 | "169.254.0.0/16", 41 | "::1/128", 42 | "fc00::/7", 43 | "fe80::/10", 44 | ]; 45 | 46 | struct SimpleLogger { 47 | local_offset: UtcOffset, 48 | } 49 | 50 | impl log::Log for SimpleLogger { 51 | fn enabled(&self, metadata: &Metadata) -> bool { 52 | metadata.level() <= Level::Trace 53 | } 54 | 55 | fn log(&self, record: &Record) { 56 | if !self.enabled(record.metadata()) { 57 | return; 58 | } 59 | 60 | let now = OffsetDateTime::now_utc().to_offset(self.local_offset); 61 | 62 | let format = format_description!( 63 | "[year]-[month]-[day] [hour]:[minute]:[second].[subsecond digits:3]" 64 | ); 65 | 66 | let mut ts = [0u8; 64]; 67 | 68 | let size = { 69 | let mut ts = io::Cursor::new(&mut ts[..]); 70 | 71 | now.format_into(&mut ts, &format) 72 | .expect("failed to write timestamp"); 73 | 74 | ts.position() as usize 75 | }; 76 | 77 | let ts = str::from_utf8(&ts[..size]).expect("timestamp is not utf-8"); 78 | 79 | let lname = match record.level() { 80 | log::Level::Error => "ERR", 81 | log::Level::Warn => "WARN", 82 | log::Level::Info => "INFO", 83 | log::Level::Debug => "DEBUG", 84 | log::Level::Trace => "TRACE", 85 | }; 86 | 87 | println!("[{}] {} [{}] {}", lname, ts, record.target(), record.args()); 88 | } 89 | 90 | fn flush(&self) {} 91 | } 92 | 93 | static mut LOGGER: mem::MaybeUninit = mem::MaybeUninit::uninit(); 94 | 95 | fn get_simple_logger() -> &'static SimpleLogger { 96 | static INIT: Once = Once::new(); 97 | 98 | unsafe { 99 | INIT.call_once(|| { 100 | let local_offset = 101 | UtcOffset::current_local_offset().expect("failed to get local time offset"); 102 | 103 | LOGGER.write(SimpleLogger { local_offset }); 104 | }); 105 | 106 | LOGGER.as_ptr().as_ref().unwrap() 107 | } 108 | } 109 | 110 | struct Args { 111 | id: String, 112 | workers: usize, 113 | req_maxconn: usize, 114 | stream_maxconn: usize, 115 | buffer_size: usize, 116 | body_buffer_size: usize, 117 | messages_max: usize, 118 | req_timeout: usize, 119 | stream_timeout: usize, 120 | listen: Vec, 121 | zclient_req_specs: Vec, 122 | zclient_stream_specs: Vec, 123 | zclient_connect: bool, 124 | zserver_req_specs: Vec, 125 | zserver_stream_specs: Vec, 126 | zserver_connect: bool, 127 | ipc_file_mode: u32, 128 | tls_identities_dir: String, 129 | allow_compression: bool, 130 | deny_out_internal: bool, 131 | } 132 | 133 | fn process_args_and_run(args: Args) -> Result<(), Box> { 134 | if args.id.is_empty() || args.id.contains(' ') { 135 | return Err("failed to parse id: value cannot be empty or contain a space".into()); 136 | } 137 | 138 | if args.workers > WORKERS_MAX { 139 | return Err("failed to parse workers: value too large".into()); 140 | } 141 | 142 | if args.req_maxconn + args.stream_maxconn > CONNS_MAX { 143 | return Err("total maxconn is too large".into()); 144 | } 145 | 146 | let mut config = app::Config { 147 | instance_id: args.id, 148 | workers: args.workers, 149 | req_maxconn: args.req_maxconn, 150 | stream_maxconn: args.stream_maxconn, 151 | buffer_size: args.buffer_size, 152 | body_buffer_size: args.body_buffer_size, 153 | messages_max: args.messages_max, 154 | req_timeout: Duration::from_secs(args.req_timeout as u64), 155 | stream_timeout: Duration::from_secs(args.stream_timeout as u64), 156 | listen: Vec::new(), 157 | zclient_req: args.zclient_req_specs, 158 | zclient_stream: args.zclient_stream_specs, 159 | zclient_connect: args.zclient_connect, 160 | zserver_req: args.zserver_req_specs, 161 | zserver_stream: args.zserver_stream_specs, 162 | zserver_connect: args.zserver_connect, 163 | ipc_file_mode: args.ipc_file_mode, 164 | certs_dir: PathBuf::from(args.tls_identities_dir), 165 | allow_compression: args.allow_compression, 166 | deny: Vec::new(), 167 | }; 168 | 169 | for v in args.listen.iter() { 170 | let mut parts = v.split(','); 171 | 172 | // there's always a first part 173 | let part1 = parts.next().unwrap(); 174 | 175 | let mut stream = true; 176 | let mut tls = false; 177 | let mut default_cert = None; 178 | let mut local = false; 179 | let mut mode = None; 180 | let mut user = None; 181 | let mut group = None; 182 | 183 | for part in parts { 184 | let (k, v) = match part.find('=') { 185 | Some(pos) => (&part[..pos], &part[(pos + 1)..]), 186 | None => (part, ""), 187 | }; 188 | 189 | match k { 190 | "req" => stream = false, 191 | "stream" => stream = true, 192 | "tls" => tls = true, 193 | "default-cert" => default_cert = Some(String::from(v)), 194 | "local" => local = true, 195 | "mode" => match u32::from_str_radix(v, 8) { 196 | Ok(x) => mode = Some(x), 197 | Err(e) => return Err(format!("failed to parse mode: {}", e).into()), 198 | }, 199 | "user" => user = Some(String::from(v)), 200 | "group" => group = Some(String::from(v)), 201 | _ => return Err(format!("failed to parse listen: invalid param: {}", part).into()), 202 | } 203 | } 204 | 205 | let spec = if local { 206 | app::ListenSpec::Local { 207 | path: PathBuf::from(part1), 208 | mode, 209 | user, 210 | group, 211 | } 212 | } else { 213 | let port_pos = match part1.rfind(':') { 214 | Some(pos) => pos + 1, 215 | None => 0, 216 | }; 217 | 218 | let port = &part1[port_pos..]; 219 | if port.parse::().is_err() { 220 | return Err(format!("failed to parse listen: invalid port {}", port).into()); 221 | } 222 | 223 | let addr = if port_pos > 0 { 224 | String::from(part1) 225 | } else { 226 | format!("0.0.0.0:{}", part1) 227 | }; 228 | 229 | let addr = match addr.parse() { 230 | Ok(addr) => addr, 231 | Err(e) => { 232 | return Err(format!("failed to parse listen: {}", e).into()); 233 | } 234 | }; 235 | 236 | app::ListenSpec::Tcp { 237 | addr, 238 | tls, 239 | default_cert, 240 | } 241 | }; 242 | 243 | config.listen.push(app::ListenConfig { spec, stream }); 244 | } 245 | 246 | if args.deny_out_internal { 247 | for s in PRIVATE_SUBNETS.iter() { 248 | config.deny.push(s.parse().unwrap()); 249 | } 250 | } 251 | 252 | condure::run(&config) 253 | } 254 | 255 | fn main() { 256 | let matches = Command::new("condure") 257 | .version(crate_version!()) 258 | .about("HTTP/WebSocket connection manager") 259 | .arg( 260 | Arg::new("log-level") 261 | .long("log-level") 262 | .num_args(1) 263 | .value_name("N") 264 | .help("Log level") 265 | .default_value("2"), 266 | ) 267 | .arg( 268 | Arg::new("id") 269 | .long("id") 270 | .num_args(1) 271 | .value_name("ID") 272 | .help("Instance ID") 273 | .default_value("condure"), 274 | ) 275 | .arg( 276 | Arg::new("workers") 277 | .long("workers") 278 | .num_args(1) 279 | .value_name("N") 280 | .help("Number of worker threads") 281 | .default_value("2"), 282 | ) 283 | .arg( 284 | Arg::new("req-maxconn") 285 | .long("req-maxconn") 286 | .num_args(1) 287 | .value_name("N") 288 | .help("Maximum number of concurrent connections in req mode") 289 | .default_value("100"), 290 | ) 291 | .arg( 292 | Arg::new("stream-maxconn") 293 | .long("stream-maxconn") 294 | .num_args(1) 295 | .value_name("N") 296 | .help("Maximum number of concurrent connections in stream mode") 297 | .default_value("10000"), 298 | ) 299 | .arg( 300 | Arg::new("buffer-size") 301 | .long("buffer-size") 302 | .num_args(1) 303 | .value_name("N") 304 | .help("Connection buffer size (two buffers per connection)") 305 | .default_value("8192"), 306 | ) 307 | .arg( 308 | Arg::new("body-buffer-size") 309 | .long("body-buffer-size") 310 | .num_args(1) 311 | .value_name("N") 312 | .help("Body buffer size for connections in req mode") 313 | .default_value("100000"), 314 | ) 315 | .arg( 316 | Arg::new("messages-max") 317 | .long("messages-max") 318 | .num_args(1) 319 | .value_name("N") 320 | .help("Maximum number of queued WebSocket messages per connection") 321 | .default_value("100"), 322 | ) 323 | .arg( 324 | Arg::new("req-timeout") 325 | .long("req-timeout") 326 | .num_args(1) 327 | .value_name("N") 328 | .help("Connection timeout in req mode (seconds)") 329 | .default_value("30"), 330 | ) 331 | .arg( 332 | Arg::new("stream-timeout") 333 | .long("stream-timeout") 334 | .num_args(1) 335 | .value_name("N") 336 | .help("Connection timeout in stream mode (seconds)") 337 | .default_value("1800"), 338 | ) 339 | .arg( 340 | Arg::new("listen") 341 | .long("listen") 342 | .num_args(1) 343 | .value_name("[addr:]port[,params...]") 344 | .action(ArgAction::Append) 345 | .help("Port to listen on"), 346 | ) 347 | .arg( 348 | Arg::new("zclient-req") 349 | .long("zclient-req") 350 | .num_args(1) 351 | .value_name("spec") 352 | .action(ArgAction::Append) 353 | .help("ZeroMQ client REQ spec") 354 | .default_value("ipc://client"), 355 | ) 356 | .arg( 357 | Arg::new("zclient-stream") 358 | .long("zclient-stream") 359 | .num_args(1) 360 | .value_name("spec-base") 361 | .action(ArgAction::Append) 362 | .help("ZeroMQ client PUSH/ROUTER/SUB spec base") 363 | .default_value("ipc://client"), 364 | ) 365 | .arg( 366 | Arg::new("zclient-connect") 367 | .long("zclient-connect") 368 | .action(ArgAction::SetTrue) 369 | .help("ZeroMQ client sockets should connect instead of bind"), 370 | ) 371 | .arg( 372 | Arg::new("zserver-req") 373 | .long("zserver-req") 374 | .num_args(1) 375 | .value_name("spec") 376 | .action(ArgAction::Append) 377 | .help("ZeroMQ server REQ spec"), 378 | ) 379 | .arg( 380 | Arg::new("zserver-stream") 381 | .long("zserver-stream") 382 | .num_args(1) 383 | .value_name("spec-base") 384 | .action(ArgAction::Append) 385 | .help("ZeroMQ server PULL/ROUTER/PUB spec base"), 386 | ) 387 | .arg( 388 | Arg::new("zserver-connect") 389 | .long("zserver-connect") 390 | .action(ArgAction::SetTrue) 391 | .help("ZeroMQ server sockets should connect instead of bind"), 392 | ) 393 | .arg( 394 | Arg::new("ipc-file-mode") 395 | .long("ipc-file-mode") 396 | .num_args(1) 397 | .value_name("octal") 398 | .help("Permissions for ZeroMQ IPC binds"), 399 | ) 400 | .arg( 401 | Arg::new("tls-identities-dir") 402 | .long("tls-identities-dir") 403 | .num_args(1) 404 | .value_name("directory") 405 | .help("Directory containing certificates and private keys") 406 | .default_value("."), 407 | ) 408 | .arg( 409 | Arg::new("compression") 410 | .long("compression") 411 | .action(ArgAction::SetTrue) 412 | .help("Allow compression to be used"), 413 | ) 414 | .arg( 415 | Arg::new("deny-out-internal") 416 | .long("deny-out-internal") 417 | .action(ArgAction::SetTrue) 418 | .help("Block outbound connections to local/internal IP address ranges"), 419 | ) 420 | .arg( 421 | Arg::new("sizes") 422 | .long("sizes") 423 | .action(ArgAction::SetTrue) 424 | .help("Prints sizes of tasks and other objects"), 425 | ) 426 | .get_matches(); 427 | 428 | log::set_logger(get_simple_logger()).unwrap(); 429 | 430 | log::set_max_level(LevelFilter::Info); 431 | 432 | let level = matches.get_one::("log-level").unwrap(); 433 | 434 | let level: usize = match level.parse() { 435 | Ok(x) => x, 436 | Err(e) => { 437 | error!("failed to parse log-level: {}", e); 438 | process::exit(1); 439 | } 440 | }; 441 | 442 | let level = match level { 443 | 0 => LevelFilter::Error, 444 | 1 => LevelFilter::Warn, 445 | 2 => LevelFilter::Info, 446 | 3 => LevelFilter::Debug, 447 | 4..=core::usize::MAX => LevelFilter::Trace, 448 | _ => unreachable!(), 449 | }; 450 | 451 | log::set_max_level(level); 452 | 453 | if *matches.get_one("sizes").unwrap() { 454 | for (name, size) in condure::app::App::sizes() { 455 | println!("{}: {} bytes", name, size); 456 | } 457 | process::exit(0); 458 | } 459 | 460 | let id = matches.get_one::("id").unwrap(); 461 | 462 | let workers = matches.get_one::("workers").unwrap(); 463 | 464 | let workers: usize = match workers.parse() { 465 | Ok(x) => x, 466 | Err(e) => { 467 | error!("failed to parse workers: {}", e); 468 | process::exit(1); 469 | } 470 | }; 471 | 472 | let req_maxconn = matches.get_one::("req-maxconn").unwrap(); 473 | 474 | let req_maxconn: usize = match req_maxconn.parse() { 475 | Ok(x) => x, 476 | Err(e) => { 477 | error!("failed to parse req-maxconn: {}", e); 478 | process::exit(1); 479 | } 480 | }; 481 | 482 | let stream_maxconn = matches.get_one::("stream-maxconn").unwrap(); 483 | 484 | let stream_maxconn: usize = match stream_maxconn.parse() { 485 | Ok(x) => x, 486 | Err(e) => { 487 | error!("failed to parse stream-maxconn: {}", e); 488 | process::exit(1); 489 | } 490 | }; 491 | 492 | let buffer_size = matches.get_one::("buffer-size").unwrap(); 493 | 494 | let buffer_size: usize = match buffer_size.parse() { 495 | Ok(x) => x, 496 | Err(e) => { 497 | error!("failed to parse buffer-size: {}", e); 498 | process::exit(1); 499 | } 500 | }; 501 | 502 | let body_buffer_size = matches.get_one::("body-buffer-size").unwrap(); 503 | 504 | let body_buffer_size: usize = match body_buffer_size.parse() { 505 | Ok(x) => x, 506 | Err(e) => { 507 | error!("failed to parse body-buffer-size: {}", e); 508 | process::exit(1); 509 | } 510 | }; 511 | 512 | let messages_max = matches.get_one::("messages-max").unwrap(); 513 | 514 | let messages_max: usize = match messages_max.parse() { 515 | Ok(x) => x, 516 | Err(e) => { 517 | error!("failed to parse messages-max: {}", e); 518 | process::exit(1); 519 | } 520 | }; 521 | 522 | let req_timeout = matches.get_one::("req-timeout").unwrap(); 523 | 524 | let req_timeout: usize = match req_timeout.parse() { 525 | Ok(x) => x, 526 | Err(e) => { 527 | error!("failed to parse req-timeout: {}", e); 528 | process::exit(1); 529 | } 530 | }; 531 | 532 | let stream_timeout = matches.get_one::("stream-timeout").unwrap(); 533 | 534 | let stream_timeout: usize = match stream_timeout.parse() { 535 | Ok(x) => x, 536 | Err(e) => { 537 | error!("failed to parse stream-timeout: {}", e); 538 | process::exit(1); 539 | } 540 | }; 541 | 542 | let mut listen: Vec = matches 543 | .get_many::("listen") 544 | .unwrap_or_default() 545 | .map(|v| v.to_owned()) 546 | .collect(); 547 | 548 | let zclient_req_specs: Vec = matches 549 | .get_many::("zclient-req") 550 | .unwrap() 551 | .map(|v| v.to_owned()) 552 | .collect(); 553 | 554 | let zclient_stream_specs: Vec = matches 555 | .get_many::("zclient-stream") 556 | .unwrap() 557 | .map(|v| v.to_owned()) 558 | .collect(); 559 | 560 | let zclient_connect = *matches.get_one("zclient-connect").unwrap(); 561 | 562 | let zserver_req_specs: Vec = matches 563 | .get_many::("zserver-req") 564 | .unwrap_or_default() 565 | .map(|v| v.to_owned()) 566 | .collect(); 567 | 568 | let zserver_stream_specs: Vec = matches 569 | .get_many::("zserver-stream") 570 | .unwrap_or_default() 571 | .map(|v| v.to_owned()) 572 | .collect(); 573 | 574 | let zserver_connect = *matches.get_one("zserver-connect").unwrap(); 575 | 576 | let ipc_file_mode = matches 577 | .get_one::("ipc-file-mode") 578 | .cloned() 579 | .unwrap_or(String::from("0")); 580 | 581 | let ipc_file_mode = match u32::from_str_radix(&ipc_file_mode, 8) { 582 | Ok(x) => x, 583 | Err(e) => { 584 | error!("failed to parse ipc-file-mode: {}", e); 585 | process::exit(1); 586 | } 587 | }; 588 | 589 | let tls_identities_dir = matches.get_one::("tls-identities-dir").unwrap(); 590 | 591 | let allow_compression = *matches.get_one("compression").unwrap(); 592 | 593 | let deny_out_internal = *matches.get_one("deny-out-internal").unwrap(); 594 | 595 | // if no zmq server specs are set (needed by client mode), specify 596 | // default listen configuration in order to enable server mode. this 597 | // means if zmq server specs are set, then server mode won't be enabled 598 | // by default 599 | if listen.is_empty() && zserver_req_specs.is_empty() && zserver_stream_specs.is_empty() { 600 | listen.push("0.0.0.0:8000,stream".to_string()); 601 | } 602 | 603 | let args = Args { 604 | id: id.to_string(), 605 | workers, 606 | req_maxconn, 607 | stream_maxconn, 608 | buffer_size, 609 | body_buffer_size, 610 | messages_max, 611 | req_timeout, 612 | stream_timeout, 613 | listen, 614 | zclient_req_specs, 615 | zclient_stream_specs, 616 | zclient_connect, 617 | zserver_req_specs, 618 | zserver_stream_specs, 619 | zserver_connect, 620 | ipc_file_mode, 621 | tls_identities_dir: tls_identities_dir.to_string(), 622 | allow_compression, 623 | deny_out_internal, 624 | }; 625 | 626 | if let Err(e) = process_args_and_run(args) { 627 | error!("{}", e); 628 | process::exit(1); 629 | } 630 | } 631 | --------------------------------------------------------------------------------