livekit_token_source/
caching.rs

1// Copyright 2026 LiveKit, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::sync::Mutex;
16
17use async_trait::async_trait;
18
19use crate::request::TokenSourceFetchOptions;
20use crate::response::{TokenSourceResponse, TokenSourceResult};
21use crate::token_source::TokenSourceConfigurable;
22
23/// Persistence backend for [`TokenSourceCached`].
24///
25/// Implement this trait to keep credentials in a custom location, e.g. the
26/// platform keychain or a database. The default is the process-lifetime
27/// [`TokenSourceInMemoryStore`].
28#[async_trait]
29pub trait TokenSourceStore: Send + Sync {
30    /// Stores the given credentials, replacing any stored previously.
31    async fn store(&self, options: TokenSourceFetchOptions, response: TokenSourceResponse);
32
33    /// Returns the stored credentials, or `None` if nothing is stored.
34    async fn retrieve(&self) -> Option<(TokenSourceFetchOptions, TokenSourceResponse)>;
35
36    /// Removes the stored credentials.
37    async fn clear(&self);
38}
39
40/// The default [`TokenSourceStore`]: keeps credentials in memory, losing them
41/// when the process exits.
42#[derive(Default)]
43pub struct TokenSourceInMemoryStore {
44    cached: Mutex<Option<(TokenSourceFetchOptions, TokenSourceResponse)>>,
45}
46
47#[async_trait]
48impl TokenSourceStore for TokenSourceInMemoryStore {
49    // Lock poisoning is recovered from rather than propagated: the guarded
50    // data is a plain `Option`, left valid even if a holder panicked mid-way.
51    async fn store(&self, options: TokenSourceFetchOptions, response: TokenSourceResponse) {
52        *self.cached.lock().unwrap_or_else(|e| e.into_inner()) = Some((options, response));
53    }
54
55    async fn retrieve(&self) -> Option<(TokenSourceFetchOptions, TokenSourceResponse)> {
56        self.cached.lock().unwrap_or_else(|e| e.into_inner()).clone()
57    }
58
59    async fn clear(&self) {
60        *self.cached.lock().unwrap_or_else(|e| e.into_inner()) = None;
61    }
62}
63
64type Validator = Box<dyn Fn(&TokenSourceFetchOptions, &TokenSourceResponse) -> bool + Send + Sync>;
65
66/// The return type of [`TokenSourceConfigurable::cached`]: wraps another token
67/// source and stores the last fetched credentials, serving them for repeat
68/// fetches with equal options for as long as they stay valid.
69///
70/// Credentials are kept in a [`TokenSourceStore`] (by default in memory, see
71/// [`TokenSourceCached::with_store`]) and are considered valid as long as a
72/// validator accepts them (by default until the token expires, see
73/// [`TokenSourceCached::with_validator`]).
74///
75/// Concurrent fetches that miss the cache are not deduplicated; each hits the
76/// underlying source and the last response wins.
77pub struct TokenSourceCached<S> {
78    source: S,
79    store: Box<dyn TokenSourceStore>,
80    validator: Validator,
81}
82
83impl<S> TokenSourceCached<S> {
84    pub(crate) fn new(source: S) -> Self {
85        Self {
86            source,
87            store: Box::new(TokenSourceInMemoryStore::default()),
88            validator: Box::new(|_, response| response.has_valid_token()),
89        }
90    }
91
92    /// Replaces the default in-memory store with a custom [`TokenSourceStore`],
93    /// e.g. one persisting credentials to the platform keychain or a database.
94    pub fn with_store(mut self, store: impl TokenSourceStore + 'static) -> Self {
95        self.store = Box::new(store);
96        self
97    }
98
99    /// Replaces the default validator deciding whether stored credentials are
100    /// still valid or must be refetched. The default checks that the token has
101    /// not expired ([`TokenSourceResponse::has_valid_token`]).
102    ///
103    /// Pass the closure inline (or annotate its parameter types); binding it
104    /// to a variable first can fail closure type inference.
105    pub fn with_validator<F>(mut self, validator: F) -> Self
106    where
107        F: Fn(&TokenSourceFetchOptions, &TokenSourceResponse) -> bool + Send + Sync + 'static,
108    {
109        self.validator = Box::new(validator);
110        self
111    }
112
113    /// Removes the stored credentials, forcing the next fetch to hit the
114    /// underlying source.
115    ///
116    /// A fetch already in flight is unaffected: it still resolves and stores
117    /// its response afterwards, repopulating the cache (last writer wins).
118    pub async fn invalidate(&self) {
119        self.store.clear().await;
120    }
121
122    /// Returns the last stored response, if any — without checking its
123    /// validity or which options it was fetched with, so it may be expired.
124    pub async fn cached_response(&self) -> Option<TokenSourceResponse> {
125        self.store.retrieve().await.map(|(_, response)| response)
126    }
127}
128
129#[async_trait]
130impl<S: TokenSourceConfigurable + Send + Sync> TokenSourceConfigurable for TokenSourceCached<S> {
131    async fn fetch(
132        &self,
133        options: &TokenSourceFetchOptions,
134    ) -> TokenSourceResult<TokenSourceResponse> {
135        if let Some((cached_options, cached_response)) = self.store.retrieve().await {
136            if cached_options == *options && (self.validator)(&cached_options, &cached_response) {
137                return Ok(cached_response);
138            }
139        }
140
141        let response = self.source.fetch(options).await?;
142        self.store.store(options.clone(), response.clone()).await;
143        Ok(response)
144    }
145}
livekit_token_source/caching.rs

livekit_token_source/
caching.rs
