pipeline/manager/
pipeline_cache.rs

1// Copyright 2023 Greptime Team
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::future::Future;
16use std::sync::Arc;
17use std::time::Duration;
18
19use datatypes::timestamp::TimestampNanosecond;
20use moka::future::Cache;
21
22use crate::error::{CacheLoadSnafu, MultiPipelineWithDiffSchemaSnafu, Result};
23use crate::etl::Pipeline;
24use crate::manager::PipelineVersion;
25use crate::table::EMPTY_SCHEMA_NAME;
26use crate::util::{generate_pipeline_cache_key, generate_pipeline_cache_key_suffix};
27
28/// Pipeline table cache size.
29const PIPELINES_CACHE_SIZE: u64 = 10000;
30
31/// Pipeline cache is located on a separate file on purpose,
32/// to encapsulate inner cache. Only public methods are exposed.
33///
34/// `pipelines` and `original_pipelines` are keyed by the *requested* schema so
35/// a lookup is a single key probe, as [`Cache::try_get_with`] requires;
36/// resolving it to a stored schema is the loader's job. `failover_cache` has no
37/// loader and keeps the stored-schema key.
38pub(crate) struct PipelineCache {
39    pipelines: Cache<String, Arc<Pipeline>>,
40    original_pipelines: Cache<String, PipelineContent>,
41    /// If the pipeline table is invalid, we can use this cache to prevent failures when writing logs through the pipeline
42    /// The failover cache never expires, but it will be updated when the pipelines cache is updated.
43    failover_cache: Cache<String, PipelineContent>,
44}
45
46#[derive(Clone, Debug, PartialEq, Eq)]
47pub struct PipelineContent {
48    pub name: String,
49    pub content: String,
50    pub version: TimestampNanosecond,
51    pub schema: String,
52}
53
54impl PipelineCache {
55    pub(crate) fn new(ttl: Duration) -> Self {
56        Self {
57            pipelines: Cache::builder()
58                .max_capacity(PIPELINES_CACHE_SIZE)
59                .time_to_live(ttl)
60                .name("pipelines")
61                .build(),
62            original_pipelines: Cache::builder()
63                .max_capacity(PIPELINES_CACHE_SIZE)
64                .time_to_live(ttl)
65                .name("original_pipelines")
66                .build(),
67            failover_cache: Cache::builder()
68                .max_capacity(PIPELINES_CACHE_SIZE)
69                .name("failover_cache")
70                .build(),
71        }
72    }
73
74    /// Concurrent misses on the same key share one `init` call.
75    pub(crate) async fn get_pipeline_with(
76        &self,
77        schema: &str,
78        name: &str,
79        version: PipelineVersion,
80        init: impl Future<Output = Result<Arc<Pipeline>>>,
81    ) -> Result<Arc<Pipeline>> {
82        let key = generate_pipeline_cache_key(schema, name, version);
83        self.pipelines
84            .try_get_with(key, init)
85            .await
86            .map_err(|error| CacheLoadSnafu { error }.build())
87    }
88
89    /// Concurrent misses on the same key share one `init` call.
90    pub(crate) async fn get_pipeline_str_with(
91        &self,
92        schema: &str,
93        name: &str,
94        version: PipelineVersion,
95        init: impl Future<Output = Result<PipelineContent>>,
96    ) -> Result<PipelineContent> {
97        let key = generate_pipeline_cache_key(schema, name, version);
98        self.original_pipelines
99            .try_get_with(key, init)
100            .await
101            .map_err(|error| CacheLoadSnafu { error }.build())
102    }
103
104    /// Resolves across schemas, unlike the loaded caches: a pipeline stored
105    /// under the empty schema is reachable from any schema.
106    pub(crate) async fn get_failover_cache(
107        &self,
108        schema: &str,
109        name: &str,
110        version: PipelineVersion,
111    ) -> Result<Option<PipelineContent>> {
112        for key in [
113            generate_pipeline_cache_key(EMPTY_SCHEMA_NAME, name, version),
114            generate_pipeline_cache_key(schema, name, version),
115        ] {
116            if let Some(content) = self.failover_cache.get(&key).await {
117                return Ok(Some(content));
118            }
119        }
120
121        // Stored under some other schema; unambiguous only if exactly one has it.
122        let suffix = generate_pipeline_cache_key_suffix(name, version);
123        let mut found = self
124            .failover_cache
125            .iter()
126            .filter(|(k, _)| k.ends_with(&suffix))
127            .collect::<Vec<_>>();
128
129        match found.len() {
130            0 => Ok(None),
131            1 => Ok(Some(found.remove(0).1)),
132            _ => MultiPipelineWithDiffSchemaSnafu {
133                name: name.to_string(),
134                current_schema: schema.to_string(),
135                schemas: found
136                    .iter()
137                    .filter_map(|(k, _)| k.split_once('/').map(|k| k.0))
138                    .collect::<Vec<_>>()
139                    .join(","),
140            }
141            .fail(),
142        }
143    }
144
145    pub(crate) async fn insert_failover_cache(&self, content: PipelineContent, with_latest: bool) {
146        let versioned =
147            generate_pipeline_cache_key(&content.schema, &content.name, Some(content.version));
148        let latest = generate_pipeline_cache_key(&content.schema, &content.name, None);
149
150        self.failover_cache.insert(versioned, content.clone()).await;
151        if with_latest {
152            self.failover_cache.insert(latest, content).await;
153        }
154    }
155
156    /// Dropping the stale `latest` aliases also clears the failover entries, so
157    /// the new version is written back: an outage before the first read-back
158    /// would otherwise have nothing to fall back on.
159    pub(crate) async fn on_pipeline_created(&self, content: PipelineContent) {
160        self.invalidate(&content.name, None).await;
161        self.insert_failover_cache(content, true).await;
162    }
163
164    /// Sweeps every schema and all three caches: the `latest` alias always,
165    /// plus `version` when given.
166    pub(crate) async fn invalidate(&self, name: &str, version: PipelineVersion) {
167        let mut suffixes = vec![generate_pipeline_cache_key_suffix(name, None)];
168        if version.is_some() {
169            suffixes.push(generate_pipeline_cache_key_suffix(name, version));
170        }
171
172        let ks = self
173            .pipelines
174            .iter()
175            .map(|(k, _)| k)
176            .chain(self.original_pipelines.iter().map(|(k, _)| k))
177            .chain(self.failover_cache.iter().map(|(k, _)| k))
178            .filter(|k| suffixes.iter().any(|suffix| k.ends_with(suffix)))
179            .collect::<Vec<_>>();
180
181        for k in ks {
182            let k = k.as_str();
183            self.pipelines.invalidate(k).await;
184            self.original_pipelines.invalidate(k).await;
185            self.failover_cache.invalidate(k).await;
186        }
187    }
188}
189
190#[cfg(test)]
191mod tests {
192    use std::sync::atomic::{AtomicUsize, Ordering};
193
194    use tokio::sync::Barrier;
195
196    use super::*;
197
198    /// Stored under the empty schema, i.e. visible from every schema.
199    fn content_at(version: i64) -> PipelineContent {
200        PipelineContent {
201            name: "p".to_string(),
202            content: "transform:".to_string(),
203            version: TimestampNanosecond::new(version),
204            schema: EMPTY_SCHEMA_NAME.to_string(),
205        }
206    }
207
208    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
209    async fn test_concurrent_misses_run_one_loader() {
210        const CONCURRENCY: usize = 8;
211
212        let cache = Arc::new(PipelineCache::new(Duration::from_secs(60)));
213        let loads = Arc::new(AtomicUsize::new(0));
214        let barrier = Arc::new(Barrier::new(CONCURRENCY));
215
216        let handles = (0..CONCURRENCY)
217            .map(|_| {
218                let (cache, loads, barrier) = (cache.clone(), loads.clone(), barrier.clone());
219                tokio::spawn(async move {
220                    barrier.wait().await;
221                    cache
222                        .get_pipeline_str_with("db", "p", None, async {
223                            loads.fetch_add(1, Ordering::SeqCst);
224                            // Hold the loader open so every caller is waiting on it.
225                            tokio::time::sleep(Duration::from_millis(100)).await;
226                            Ok(content_at(1))
227                        })
228                        .await
229                        .unwrap()
230                })
231            })
232            .collect::<Vec<_>>();
233
234        for handle in handles {
235            assert_eq!(handle.await.unwrap(), content_at(1));
236        }
237        assert_eq!(loads.load(Ordering::SeqCst), 1);
238    }
239
240    #[tokio::test]
241    async fn test_delete_drops_version_pinned_entry() {
242        let cache = PipelineCache::new(Duration::from_secs(60));
243        let content = content_at(1);
244        let version = Some(content.version);
245
246        cache
247            .get_pipeline_str_with("db", "p", version, async { Ok(content.clone()) })
248            .await
249            .unwrap();
250
251        cache.invalidate("p", version).await;
252
253        let loads = AtomicUsize::new(0);
254        cache
255            .get_pipeline_str_with("db", "p", version, async {
256                loads.fetch_add(1, Ordering::SeqCst);
257                Ok(content.clone())
258            })
259            .await
260            .unwrap();
261        assert_eq!(loads.load(Ordering::SeqCst), 1);
262    }
263
264    #[tokio::test]
265    async fn test_create_drops_stale_latest_and_primes_failover() {
266        let cache = PipelineCache::new(Duration::from_secs(60));
267        let v2 = content_at(2);
268
269        cache
270            .get_pipeline_str_with("a", "p", None, async { Ok(content_at(1)) })
271            .await
272            .unwrap();
273        cache.insert_failover_cache(content_at(1), true).await;
274
275        cache.on_pipeline_created(v2.clone()).await;
276
277        let loads = AtomicUsize::new(0);
278        let cached = cache
279            .get_pipeline_str_with("a", "p", None, async {
280                loads.fetch_add(1, Ordering::SeqCst);
281                Ok(v2.clone())
282            })
283            .await
284            .unwrap();
285        assert_eq!(loads.load(Ordering::SeqCst), 1);
286        assert_eq!(cached.version, v2.version);
287
288        let failover = cache.get_failover_cache("b", "p", None).await.unwrap();
289        assert_eq!(failover.map(|c| c.version), Some(v2.version));
290    }
291
292    #[tokio::test]
293    async fn test_failover_serves_global_pipeline_to_unwarmed_schema() {
294        let cache = PipelineCache::new(Duration::from_secs(60));
295        let content = content_at(1);
296
297        cache.insert_failover_cache(content.clone(), true).await;
298
299        let found = cache.get_failover_cache("b", "p", None).await.unwrap();
300        assert_eq!(found, Some(content.clone()));
301
302        // A same-named pipeline under another schema must not shadow the global one.
303        let schema_local = PipelineContent {
304            schema: "x".to_string(),
305            ..content_at(2)
306        };
307        cache.insert_failover_cache(schema_local, true).await;
308
309        let found = cache.get_failover_cache("b", "p", None).await.unwrap();
310        assert_eq!(found, Some(content));
311    }
312}
pipeline/manager/pipeline_cache.rs

pipeline/manager/
pipeline_cache.rs
