tinyvector 0.0.1

A tiny nearest-neighbor embedding database built with SQLite.
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109

use axum::Extension;
use std::{
	error::Error,
	fmt,
	fmt::Display,
	future::Future,
	sync::atomic::{AtomicBool, Ordering},
};
use tokio::{signal, sync::broadcast};

#[derive(Debug, PartialEq, Eq)]
pub struct AlreadyCreatedError;

impl Error for AlreadyCreatedError {}

impl Display for AlreadyCreatedError {
	fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
		f.write_str("shutdown handler already created")
	}
}

static CREATED: AtomicBool = AtomicBool::new(false);

#[derive(Debug, Clone)]
pub struct Shutdown {
	pub sender: broadcast::Sender<()>,
}

#[derive(Debug, Clone)]
pub struct Agent {
	sender: broadcast::Sender<()>,
}

impl Agent {
	pub fn start(&self) {
		self.sender.send(()).ok();
	}
}

impl Shutdown {
	pub fn new() -> Result<Self, AlreadyCreatedError> {
		if (CREATED).swap(true, Ordering::SeqCst) {
			tracing::error!("shutdown handler called twice");
			return Err(AlreadyCreatedError);
		}

		let (tx, _) = broadcast::channel(1);
		let handle = register_handlers();

		let tx_for_handle = tx.clone();
		tokio::spawn(async move {
			tracing::debug!("Registered shutdown handlers");
			handle.await;
			tx_for_handle.send(()).ok();
		});

		Ok(Self { sender: tx })
	}

	pub fn handle(&self) -> impl Future<Output = ()> + '_ {
		let mut rx = self.sender.subscribe();

		async move {
			let rx = rx.recv();

			rx.await.unwrap();
		}
	}

	pub fn agent(&self) -> Agent {
		Agent {
			sender: self.sender.clone(),
		}
	}

	pub fn extension(&self) -> Extension<Agent> {
		Extension(self.agent())
	}
}

fn register_handlers() -> impl Future<Output = ()> {
	let ctrl_c = async {
		signal::ctrl_c()
			.await
			.expect("failed to install Ctrl+C handler");
	};

	#[cfg(unix)]
	let terminate = async {
		signal::unix::signal(signal::unix::SignalKind::terminate())
			.expect("failed to install signal handler")
			.recv()
			.await;
	};

	#[cfg(not(unix))]
	let terminate = std::future::pending::<()>();

	async move {
		tokio::select! {
			_ = ctrl_c => {
				tracing::info!("Received Ctrl+C signal");
			},
			_ = terminate => {
				tracing::info!("Received terminate signal");
			},
		}
	}
}