9450933460

Committed 10 Jun 2024 02:45PM UTC coverage: 90.594% (+0.01%) from 90.584%

Build # 9450933460

Build Type

push

github

Committed by

web-flow

Commit Message

Merge pull request #965 from glombiewski/oidc-login-user-infos

Missing attributes in session results

Run Details

3 of 3 new or added lines in 1 file covered. (100.0%)

8 existing lines in 5 files now uncovered.

131461 of 145110 relevant lines covered (90.59%)

53310.47 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

94.56

/services/src/pro/quota/mod.rs

use std::{collections::HashMap, time::Duration};

use geoengine_operators::pro::meta::quota::{
    ComputationContext, ComputationUnit, QuotaMessage, QuotaTracking,
};
use snafu::Snafu;

use tokio::sync::mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender};

use geoengine_datatypes::util::test::TestDefault;

use crate::pro::users::UserId;

use super::{
    users::{UserDb, UserSession},
    util::config::QuotaTrackingMode,
};

#[derive(Debug, Snafu)]
#[snafu(visibility(pub(crate)))]
#[snafu(context(suffix(false)))]
pub enum QuotaError {
    QuotaExhausted,
}

#[derive(Debug, Clone)]
pub struct QuotaTrackingFactory {
    quota_sender: UnboundedSender<QuotaMessage>,
}

impl QuotaTrackingFactory {
    pub fn new(quota_sender: UnboundedSender<QuotaMessage>) -> Self {
        Self { quota_sender }
    }

    pub fn create_quota_tracking(
        &self,
        session: &UserSession,
        context: ComputationContext,
    ) -> QuotaTracking {
        QuotaTracking::new(
            self.quota_sender.clone(),
            ComputationUnit {
                issuer: session.user.id.0,
                context,
            },
        )
    }
}

impl TestDefault for QuotaTrackingFactory {
    fn test_default() -> Self {
        let (quota_sender, _quota_receiver) = unbounded_channel();
        Self::new(quota_sender)
    }
}

pub struct QuotaManager<U: UserDb + 'static> {
    mode: QuotaTrackingMode,
    user_db: U,
    quota_receiver: UnboundedReceiver<QuotaMessage>,
    quota_sender: UnboundedSender<QuotaMessage>,
    increment_quota_buffer: HashMap<UserId, u64>,
    increment_quota_buffer_size: usize,
    increment_quota_buffer_used: usize,
    increment_quota_buffer_timeout_seconds: u64,
}

impl<U: UserDb + 'static> QuotaManager<U> {
    pub fn new(
        mode: QuotaTrackingMode,
        user_db: U,
        quota_receiver: UnboundedReceiver<QuotaMessage>,
        quota_sender: UnboundedSender<QuotaMessage>,
        increment_quota_buffer_size: usize,
        increment_quota_buffer_timeout_seconds: u64,
    ) -> Self {
        Self {
            mode,
            user_db,
            quota_receiver,
            quota_sender,
            increment_quota_buffer: HashMap::new(),
            increment_quota_buffer_size,
            increment_quota_buffer_timeout_seconds,
            increment_quota_buffer_used: 0,
        }
    }

    pub fn run(mut self) {
        if self.mode == QuotaTrackingMode::Disabled {
            // if the quota tracking is disabled, we still consume the messages in order to empty the channel, but ignore the messages
            crate::util::spawn(async move {
                while let Some(_message) = self.quota_receiver.recv().await {
                    continue;
                }
            });

            return;
        }

        crate::util::spawn(async move {
            while let Some(message) = self.quota_receiver.recv().await {
                let flush_buffer = match message {
                    QuotaMessage::ComputationUnit(computation) => {
                        // TODO: issue a tracing event instead?
                        // TODO: also log workflow, or connect the computation context to a workflow beforehand.
                        //       However, currently it is possible to reuse a context for multiple queries.
                        //       Also: the operators crate knows nothing about workflows as of yet.
                        log::trace!(
                            "Quota received. User: {}, Context: {}",
                            computation.issuer,
                            computation.context
                        );

                        let user = UserId(computation.issuer);

                        *self.increment_quota_buffer.entry(user).or_default() += 1;
                        self.increment_quota_buffer_used += 1;

                        self.increment_quota_buffer_used >= self.increment_quota_buffer_size
                    }
                    QuotaMessage::Flush => {
                        log::trace!("Flush `increment_quota_buffer'");
                        true
                    }
                };

                if flush_buffer {
                    // TODO: what to do if this fails (quota can't be recorded)? Try again later?

                    if let Err(error) = self
                        .user_db
                        .bulk_increment_quota_used(self.increment_quota_buffer.drain())
                        .await
                    {
                        log::error!("Could not increment quota for users {error:?}");
                    }

                    self.increment_quota_buffer_used = 0;
                }
            }
        });

        // Periodically flush the increment_quota_buffer
        crate::util::spawn(async move {
            let timeout_duration = Duration::from_secs(self.increment_quota_buffer_timeout_seconds);

            loop {
                tokio::time::sleep(timeout_duration).await;
                let _ = self.quota_sender.send(QuotaMessage::Flush);
            }
        });
    }
}

pub fn initialize_quota_tracking<U: UserDb + 'static>(
    mode: QuotaTrackingMode,
    user_db: U,
    increment_quota_buffer_size: usize,
    increment_quota_buffer_timeout_seconds: u64,
) -> QuotaTrackingFactory {
    let (quota_sender, quota_receiver) = unbounded_channel::<QuotaMessage>();

    QuotaManager::new(
        mode,
        user_db,
        quota_receiver,
        quota_sender.clone(),
        increment_quota_buffer_size,
        increment_quota_buffer_timeout_seconds,
    )
    .run();

    QuotaTrackingFactory::new(quota_sender)
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::{
        contexts::{ApplicationContext, SessionContext},
        pro::{
            contexts::ProPostgresContext,
            ge_context,
            users::{UserAuth, UserCredentials, UserRegistration},
            util::tests::admin_login,
        },
    };
    use geoengine_datatypes::util::Identifier;
    use tokio_postgres::NoTls;

    #[ge_context::test]
    async fn it_tracks_quota(app_ctx: ProPostgresContext<NoTls>) {
        let _user = app_ctx
            .register_user(UserRegistration {
                email: "foo@example.com".to_string(),
                password: "secret1234".to_string(),
                real_name: "Foo Bar".to_string(),
            })
            .await
            .unwrap();

        let session = app_ctx
            .login(UserCredentials {
                email: "foo@example.com".to_string(),
                password: "secret1234".to_string(),
            })
            .await
            .unwrap();

        let admin_session = admin_login(&app_ctx).await;

        let quota = initialize_quota_tracking(
            QuotaTrackingMode::Check,
            app_ctx.session_context(admin_session.clone()).db(),
            0,
            60,
        );

        let tracking = quota.create_quota_tracking(&session, ComputationContext::new());

        tracking.work_unit_done();
        tracking.work_unit_done();

        let db = app_ctx.session_context(session).db();

        // wait for quota to be recorded
        let mut success = false;
        for _ in 0..10 {
            let used = db.quota_used().await.unwrap();
            let available = db.quota_available().await.unwrap();

            tokio::time::sleep(std::time::Duration::from_millis(100)).await;

            if used == 2 && available == 9997 {
                success = true;
                break;
            }
        }

        assert!(success);
    }

    #[ge_context::test]
    async fn it_tracks_quota_buffered(app_ctx: ProPostgresContext<NoTls>) {
        let _user = app_ctx
            .register_user(UserRegistration {
                email: "foo@example.com".to_string(),
                password: "secret1234".to_string(),
                real_name: "Foo Bar".to_string(),
            })
            .await
            .unwrap();

        let session = app_ctx
            .login(UserCredentials {
                email: "foo@example.com".to_string(),
                password: "secret1234".to_string(),
            })
            .await
            .unwrap();

        let admin_session = admin_login(&app_ctx).await;

        let quota = initialize_quota_tracking(
            QuotaTrackingMode::Check,
            app_ctx.session_context(admin_session.clone()).db(),
            5,
            60,
        );

        let tracking = quota.create_quota_tracking(&session, ComputationContext::new());

        tracking.work_unit_done();
        tracking.work_unit_done();
        tracking.work_unit_done();
        tracking.work_unit_done();
        tracking.work_unit_done();

        let db = app_ctx.session_context(session).db();

        // wait for quota to be recorded
        let mut success = false;
        for _ in 0..10 {
            let used = db.quota_used().await.unwrap();
            let available = db.quota_available().await.unwrap();

            tokio::time::sleep(std::time::Duration::from_millis(100)).await;

            if used == 5 && available == 9994 {
                success = true;
                break;
            }
        }

        assert!(success);
    }

    #[ge_context::test]
    async fn it_tracks_quota_buffered_timeout(app_ctx: ProPostgresContext<NoTls>) {
        let _user = app_ctx
            .register_user(UserRegistration {
                email: "foo@example.com".to_string(),
                password: "secret1234".to_string(),
                real_name: "Foo Bar".to_string(),
            })
            .await
            .unwrap();

        let session = app_ctx
            .login(UserCredentials {
                email: "foo@example.com".to_string(),
                password: "secret1234".to_string(),
            })
            .await
            .unwrap();

        let admin_session = admin_login(&app_ctx).await;

        let quota = initialize_quota_tracking(
            QuotaTrackingMode::Check,
            app_ctx.session_context(admin_session.clone()).db(),
            5,
            2,
        );

        let tracking = quota.create_quota_tracking(&session, ComputationContext::new());

        tracking.work_unit_done();
        tracking.work_unit_done();

        let db = app_ctx.session_context(session).db();

        tokio::time::sleep(std::time::Duration::from_secs(4)).await;

        let used = db.quota_used().await.unwrap();
        let available = db.quota_available().await.unwrap();

        assert_eq!(used, 2);
        assert_eq!(available, 9997);
    }
}

1	use std::{collections::HashMap, time::Duration};
2
3	use geoengine_operators::pro::meta::quota::{
4	ComputationContext, ComputationUnit, QuotaMessage, QuotaTracking,
5	};
6	use snafu::Snafu;
7
8	use tokio::sync::mpsc::{unbounded_channel, UnboundedReceiver, UnboundedSender};
9
10	use geoengine_datatypes::util::test::TestDefault;
11
12	use crate::pro::users::UserId;
13
14	use super::{
15	users::{UserDb, UserSession},
16	util::config::QuotaTrackingMode,
17	};
18
19	#[derive(Debug, Snafu)]	1✔
20	#[snafu(visibility(pub(crate)))]
21	#[snafu(context(suffix(false)))]
22	pub enum QuotaError {
23	QuotaExhausted,
24	}
25
26	#[derive(Debug, Clone)]	250✔
27	pub struct QuotaTrackingFactory {
28	quota_sender: UnboundedSender<QuotaMessage>,
29	}
30
31	impl QuotaTrackingFactory {
32	pub fn new(quota_sender: UnboundedSender<QuotaMessage>) -> Self {	97✔
33	Self { quota_sender }	97✔
34	}	97✔
35
36	pub fn create_quota_tracking(	11✔
37	&self,	11✔
38	session: &UserSession,	11✔
39	context: ComputationContext,	11✔
40	) -> QuotaTracking {	11✔
41	QuotaTracking::new(	11✔
42	self.quota_sender.clone(),	11✔
43	ComputationUnit {	11✔
44	issuer: session.user.id.0,	11✔
45	context,	11✔
46	},	11✔
47	)	11✔
48	}	11✔
49	}
50
51	impl TestDefault for QuotaTrackingFactory {
52	fn test_default() -> Self {	×
53	let (quota_sender, _quota_receiver) = unbounded_channel();	×
54	Self::new(quota_sender)	×
55	}	×
56	}
57
58	pub struct QuotaManager<U: UserDb + 'static> {
59	mode: QuotaTrackingMode,
60	user_db: U,
61	quota_receiver: UnboundedReceiver<QuotaMessage>,
62	quota_sender: UnboundedSender<QuotaMessage>,
63	increment_quota_buffer: HashMap<UserId, u64>,
64	increment_quota_buffer_size: usize,
65	increment_quota_buffer_used: usize,
66	increment_quota_buffer_timeout_seconds: u64,
67	}
68
69	impl<U: UserDb + 'static> QuotaManager<U> {
70	pub fn new(	97✔
71	mode: QuotaTrackingMode,	97✔
72	user_db: U,	97✔
73	quota_receiver: UnboundedReceiver<QuotaMessage>,	97✔
74	quota_sender: UnboundedSender<QuotaMessage>,	97✔
75	increment_quota_buffer_size: usize,	97✔
76	increment_quota_buffer_timeout_seconds: u64,	97✔
77	) -> Self {	97✔
78	Self {	97✔
79	mode,	97✔
80	user_db,	97✔
81	quota_receiver,	97✔
82	quota_sender,	97✔
83	increment_quota_buffer: HashMap::new(),	97✔
84	increment_quota_buffer_size,	97✔
85	increment_quota_buffer_timeout_seconds,	97✔
86	increment_quota_buffer_used: 0,	97✔
87	}	97✔
88	}	97✔
89
90	pub fn run(mut self) {	97✔
91	if self.mode == QuotaTrackingMode::Disabled {	97✔
92	// if the quota tracking is disabled, we still consume the messages in order to empty the channel, but ignore the messages
93	crate::util::spawn(async move {	×
94	while let Some(_message) = self.quota_receiver.recv().await {	×
95	continue;	×
96	}
97	});	×
98		×
99	return;	×
100	}	97✔
101		97✔
102	crate::util::spawn(async move {	97✔
103	while let Some(message) = self.quota_receiver.recv().await {	7,916✔
104	let flush_buffer = match message {	7,821✔
105	QuotaMessage::ComputationUnit(computation) => {	7,818✔
106	// TODO: issue a tracing event instead?	7,818✔
107	// TODO: also log workflow, or connect the computation context to a workflow beforehand.	7,818✔
108	// However, currently it is possible to reuse a context for multiple queries.	7,818✔
109	// Also: the operators crate knows nothing about workflows as of yet.	7,818✔
110	log::trace!(	7,818✔
111	"Quota received. User: {}, Context: {}",	×
112	computation.issuer,
113	computation.context
114	);
115
116	let user = UserId(computation.issuer);	7,818✔
117		7,818✔
118	*self.increment_quota_buffer.entry(user).or_default() += 1;	7,818✔
119	self.increment_quota_buffer_used += 1;	7,818✔
120		7,818✔
121	self.increment_quota_buffer_used >= self.increment_quota_buffer_size	7,818✔
122	}
123	QuotaMessage::Flush => {
124	log::trace!("Flush `increment_quota_buffer'");	3✔
125	true	3✔
126	}
127	};
128
129	if flush_buffer {	7,821✔
130	// TODO: what to do if this fails (quota can't be recorded)? Try again later?
131
132	if let Err(error) = self	7,815✔
133	.user_db	7,815✔
134	.bulk_increment_quota_used(self.increment_quota_buffer.drain())	7,815✔
135	.await	23,430✔
136	{
137	log::error!("Could not increment quota for users {error:?}");	16✔
138	}	7,797✔
139
140	self.increment_quota_buffer_used = 0;	7,813✔
141	}	6✔
142	}
143	});	97✔
144		97✔
145	// Periodically flush the increment_quota_buffer	97✔
146	crate::util::spawn(async move {	97✔
147	let timeout_duration = Duration::from_secs(self.increment_quota_buffer_timeout_seconds);	97✔
148
149	loop {
150	tokio::time::sleep(timeout_duration).await;	100✔
151	let _ = self.quota_sender.send(QuotaMessage::Flush);	3✔
152	}
153	});	97✔
154	}	97✔
155	}
156
157	pub fn initialize_quota_tracking<U: UserDb + 'static>(	97✔
158	mode: QuotaTrackingMode,	97✔
159	user_db: U,	97✔
160	increment_quota_buffer_size: usize,	97✔
161	increment_quota_buffer_timeout_seconds: u64,	97✔
162	) -> QuotaTrackingFactory {	97✔
163	let (quota_sender, quota_receiver) = unbounded_channel::<QuotaMessage>();	97✔
164		97✔
165	QuotaManager::new(	97✔
166	mode,	97✔
167	user_db,	97✔
168	quota_receiver,	97✔
169	quota_sender.clone(),	97✔
170	increment_quota_buffer_size,	97✔
171	increment_quota_buffer_timeout_seconds,	97✔
172	)	97✔
173	.run();	97✔
174		97✔
175	QuotaTrackingFactory::new(quota_sender)	97✔
176	}	97✔
177
178	#[cfg(test)]
179	mod tests {
180	use super::*;
181	use crate::{
182	contexts::{ApplicationContext, SessionContext},
183	pro::{
184	contexts::ProPostgresContext,
185	ge_context,
186	users::{UserAuth, UserCredentials, UserRegistration},
187	util::tests::admin_login,
188	},
189	};
190	use geoengine_datatypes::util::Identifier;
191	use tokio_postgres::NoTls;
192
193	#[ge_context::test]	3✔
194	async fn it_tracks_quota(app_ctx: ProPostgresContext<NoTls>) {	1✔
195	let _user = app_ctx	1✔
196	.register_user(UserRegistration {	1✔
197	email: "foo@example.com".to_string(),	1✔
198	password: "secret1234".to_string(),	1✔
199	real_name: "Foo Bar".to_string(),	1✔
200	})	1✔
201	.await	9✔
202	.unwrap();	1✔
203
204	let session = app_ctx	1✔
205	.login(UserCredentials {	1✔
206	email: "foo@example.com".to_string(),	1✔
207	password: "secret1234".to_string(),	1✔
208	})	1✔
209	.await	1✔
210	.unwrap();	1✔
211
212	let admin_session = admin_login(&app_ctx).await;	1✔
213
214	let quota = initialize_quota_tracking(	1✔
215	QuotaTrackingMode::Check,	1✔
216	app_ctx.session_context(admin_session.clone()).db(),	1✔
217	0,	1✔
218	60,	1✔
219	);	1✔
220		1✔
221	let tracking = quota.create_quota_tracking(&session, ComputationContext::new());	1✔
222		1✔
223	tracking.work_unit_done();	1✔
224	tracking.work_unit_done();	1✔
225		1✔
226	let db = app_ctx.session_context(session).db();	1✔
227		1✔
228	// wait for quota to be recorded	1✔
229	let mut success = false;	1✔
230	for _ in 0..10 {	1✔
231	let used = db.quota_used().await.unwrap();	3✔
232	let available = db.quota_available().await.unwrap();	3✔
233		1✔
234	tokio::time::sleep(std::time::Duration::from_millis(100)).await;	1✔
235
236	if used == 2 && available == 9997 {	1✔
237	success = true;	1✔
238	break;	1✔
UNCOV 239	}	×
240	}
241
242	assert!(success);	1✔
243	}	1✔
244
245	#[ge_context::test]	3✔
246	async fn it_tracks_quota_buffered(app_ctx: ProPostgresContext<NoTls>) {	1✔
247	let _user = app_ctx	1✔
248	.register_user(UserRegistration {	1✔
249	email: "foo@example.com".to_string(),	1✔
250	password: "secret1234".to_string(),	1✔
251	real_name: "Foo Bar".to_string(),	1✔
252	})	1✔
253	.await	11✔
254	.unwrap();	1✔
255
256	let session = app_ctx	1✔
257	.login(UserCredentials {	1✔
258	email: "foo@example.com".to_string(),	1✔
259	password: "secret1234".to_string(),	1✔
260	})	1✔
261	.await	4✔
262	.unwrap();	1✔
263
264	let admin_session = admin_login(&app_ctx).await;	1✔
265
266	let quota = initialize_quota_tracking(	1✔
267	QuotaTrackingMode::Check,	1✔
268	app_ctx.session_context(admin_session.clone()).db(),	1✔
269	5,	1✔
270	60,	1✔
271	);	1✔
272		1✔
273	let tracking = quota.create_quota_tracking(&session, ComputationContext::new());	1✔
274		1✔
275	tracking.work_unit_done();	1✔
276	tracking.work_unit_done();	1✔
277	tracking.work_unit_done();	1✔
278	tracking.work_unit_done();	1✔
279	tracking.work_unit_done();	1✔
280		1✔
281	let db = app_ctx.session_context(session).db();	1✔
282		1✔
283	// wait for quota to be recorded	1✔
284	let mut success = false;	1✔
285	for _ in 0..10 {	1✔
286	let used = db.quota_used().await.unwrap();	3✔
287	let available = db.quota_available().await.unwrap();	3✔
288		1✔
289	tokio::time::sleep(std::time::Duration::from_millis(100)).await;	1✔
290
291	if used == 5 && available == 9994 {	1✔
292	success = true;	1✔
293	break;	1✔
UNCOV 294	}	×
295	}
296
297	assert!(success);	1✔
298	}	1✔
299
300	#[ge_context::test]	3✔
301	async fn it_tracks_quota_buffered_timeout(app_ctx: ProPostgresContext<NoTls>) {	1✔
302	let _user = app_ctx	1✔
303	.register_user(UserRegistration {	1✔
304	email: "foo@example.com".to_string(),	1✔
305	password: "secret1234".to_string(),	1✔
306	real_name: "Foo Bar".to_string(),	1✔
307	})	1✔
308	.await	4✔
309	.unwrap();	1✔
310
311	let session = app_ctx	1✔
312	.login(UserCredentials {	1✔
313	email: "foo@example.com".to_string(),	1✔
314	password: "secret1234".to_string(),	1✔
315	})	1✔
316	.await	7✔
317	.unwrap();	1✔
318
319	let admin_session = admin_login(&app_ctx).await;	10✔
320
321	let quota = initialize_quota_tracking(	1✔
322	QuotaTrackingMode::Check,	1✔
323	app_ctx.session_context(admin_session.clone()).db(),	1✔
324	5,	1✔
325	2,	1✔
326	);	1✔
327		1✔
328	let tracking = quota.create_quota_tracking(&session, ComputationContext::new());	1✔
329		1✔
330	tracking.work_unit_done();	1✔
331	tracking.work_unit_done();	1✔
332		1✔
333	let db = app_ctx.session_context(session).db();	1✔
334		1✔
335	tokio::time::sleep(std::time::Duration::from_secs(4)).await;	1✔
336
337	let used = db.quota_used().await.unwrap();	3✔
338	let available = db.quota_available().await.unwrap();	3✔
339		1✔
340	assert_eq!(used, 2);	1✔
341	assert_eq!(available, 9997);	1✔
342	}	1✔
343	}

geo-engine / geoengine / 9450933460

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous