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 common_frontend::metrics;
20use datatypes::timestamp::TimestampNanosecond;
21use moka::future::Cache;
22
23use crate::error::{CacheLoadSnafu, MultiPipelineWithDiffSchemaSnafu, Result};
24use crate::etl::Pipeline;
25use crate::manager::PipelineVersion;
26use crate::table::EMPTY_SCHEMA_NAME;
27use crate::util::{generate_pipeline_cache_key, generate_pipeline_cache_key_suffix};
28
29/// Pipeline table cache size.
30const PIPELINES_CACHE_SIZE: u64 = 10000;
31
32/// Pipeline cache is located on a separate file on purpose,
33/// to encapsulate inner cache. Only public methods are exposed.
34///
35/// `pipelines` and `original_pipelines` are keyed by the *requested* schema so
36/// a lookup is a single key probe through [`Cache::entry`];
37/// resolving it to a stored schema is the loader's job. `failover_cache` has no
38/// loader and keeps the stored-schema key.
39pub(crate) struct PipelineCache {
40    pipelines: Cache<String, Arc<Pipeline>>,
41    original_pipelines: Cache<String, PipelineContent>,
42    /// If the pipeline table is invalid, we can use this cache to prevent failures when writing logs through the pipeline
43    /// The failover cache never expires, but it will be updated when the pipelines cache is updated.
44    failover_cache: Cache<String, PipelineContent>,
45}
46
47#[derive(Clone, Debug, PartialEq, Eq)]
48pub struct PipelineContent {
49    pub name: String,
50    pub content: String,
51    pub version: TimestampNanosecond,
52    pub schema: String,
53}
54
55impl PipelineCache {
56    pub(crate) fn new(ttl: Duration) -> Self {
57        Self {
58            pipelines: Cache::builder()
59                .max_capacity(PIPELINES_CACHE_SIZE)
60                .time_to_live(ttl)
61                .name("pipelines")
62                .build(),
63            original_pipelines: Cache::builder()
64                .max_capacity(PIPELINES_CACHE_SIZE)
65                .time_to_live(ttl)
66                .name("original_pipelines")
67                .build(),
68            failover_cache: Cache::builder()
69                .max_capacity(PIPELINES_CACHE_SIZE)
70                .name("failover_cache")
71                .build(),
72        }
73    }
74
75    /// Concurrent misses share one `init` call; successful waiters count as hits.
76    pub(crate) async fn get_pipeline_with(
77        &self,
78        schema: &str,
79        name: &str,
80        version: PipelineVersion,
81        init: impl Future<Output = Result<Arc<Pipeline>>>,
82    ) -> Result<Arc<Pipeline>> {
83        let key = generate_pipeline_cache_key(schema, name, version);
84        let entry = self.pipelines.entry(key).or_try_insert_with(init).await;
85        metrics::record_cache_lookup(
86            "pipeline",
87            entry.as_ref().is_ok_and(|entry| !entry.is_fresh()),
88        );
89        entry
90            .map(|entry| entry.into_value())
91            .map_err(|error| CacheLoadSnafu { error }.build())
92    }
93
94    /// Concurrent misses share one `init` call; successful waiters count as hits.
95    pub(crate) async fn get_pipeline_str_with(
96        &self,
97        schema: &str,
98        name: &str,
99        version: PipelineVersion,
100        init: impl Future<Output = Result<PipelineContent>>,
101    ) -> Result<PipelineContent> {
102        let key = generate_pipeline_cache_key(schema, name, version);
103        let entry = self
104            .original_pipelines
105            .entry(key)
106            .or_try_insert_with(init)
107            .await;
108        metrics::record_cache_lookup(
109            "pipeline_original",
110            entry.as_ref().is_ok_and(|entry| !entry.is_fresh()),
111        );
112        entry
113            .map(|entry| entry.into_value())
114            .map_err(|error| CacheLoadSnafu { error }.build())
115    }
116
117    /// Resolves across schemas, unlike the loaded caches: a pipeline stored
118    /// under the empty schema is reachable from any schema.
119    pub(crate) async fn get_failover_cache(
120        &self,
121        schema: &str,
122        name: &str,
123        version: PipelineVersion,
124    ) -> Result<Option<PipelineContent>> {
125        for key in [
126            generate_pipeline_cache_key(EMPTY_SCHEMA_NAME, name, version),
127            generate_pipeline_cache_key(schema, name, version),
128        ] {
129            if let Some(content) = self.failover_cache.get(&key).await {
130                metrics::record_cache_lookup("pipeline_failover", true);
131                return Ok(Some(content));
132            }
133        }
134
135        // Stored under some other schema; unambiguous only if exactly one has it.
136        let suffix = generate_pipeline_cache_key_suffix(name, version);
137        let mut found = self
138            .failover_cache
139            .iter()
140            .filter(|(k, _)| k.ends_with(&suffix))
141            .collect::<Vec<_>>();
142
143        let result = match found.len() {
144            0 => Ok(None),
145            1 => Ok(Some(found.remove(0).1)),
146            _ => MultiPipelineWithDiffSchemaSnafu {
147                name: name.to_string(),
148                current_schema: schema.to_string(),
149                schemas: found
150                    .iter()
151                    .filter_map(|(k, _)| k.split_once('/').map(|k| k.0))
152                    .collect::<Vec<_>>()
153                    .join(","),
154            }
155            .fail(),
156        };
157        metrics::record_cache_lookup(
158            "pipeline_failover",
159            result.as_ref().is_ok_and(|content| content.is_some()),
160        );
161        result
162    }
163
164    pub(crate) async fn insert_failover_cache(&self, content: PipelineContent, with_latest: bool) {
165        let versioned =
166            generate_pipeline_cache_key(&content.schema, &content.name, Some(content.version));
167        let latest = generate_pipeline_cache_key(&content.schema, &content.name, None);
168
169        self.failover_cache.insert(versioned, content.clone()).await;
170        if with_latest {
171            self.failover_cache.insert(latest, content).await;
172        }
173    }
174
175    /// Dropping the stale `latest` aliases also clears the failover entries, so
176    /// the new version is written back: an outage before the first read-back
177    /// would otherwise have nothing to fall back on.
178    pub(crate) async fn on_pipeline_created(&self, content: PipelineContent) {
179        self.invalidate(&content.name, None).await;
180        self.insert_failover_cache(content, true).await;
181    }
182
183    /// Sweeps every schema and all three caches: the `latest` alias always,
184    /// plus `version` when given.
185    pub(crate) async fn invalidate(&self, name: &str, version: PipelineVersion) {
186        let mut suffixes = vec![generate_pipeline_cache_key_suffix(name, None)];
187        if version.is_some() {
188            suffixes.push(generate_pipeline_cache_key_suffix(name, version));
189        }
190
191        let ks = self
192            .pipelines
193            .iter()
194            .map(|(k, _)| k)
195            .chain(self.original_pipelines.iter().map(|(k, _)| k))
196            .chain(self.failover_cache.iter().map(|(k, _)| k))
197            .filter(|k| suffixes.iter().any(|suffix| k.ends_with(suffix)))
198            .collect::<Vec<_>>();
199
200        for k in ks {
201            let k = k.as_str();
202            self.pipelines.invalidate(k).await;
203            self.original_pipelines.invalidate(k).await;
204            self.failover_cache.invalidate(k).await;
205        }
206    }
207}
208
209#[cfg(test)]
210mod tests {
211    use std::sync::atomic::{AtomicUsize, Ordering};
212
213    use common_frontend::metrics::{CACHE_HIT, CACHE_MISS};
214    use tokio::sync::Notify;
215
216    use super::*;
217
218    /// Stored under the empty schema, i.e. visible from every schema.
219    fn content_at(version: i64) -> PipelineContent {
220        PipelineContent {
221            name: "p".to_string(),
222            content: "transform:".to_string(),
223            version: TimestampNanosecond::new(version),
224            schema: EMPTY_SCHEMA_NAME.to_string(),
225        }
226    }
227
228    #[tokio::test]
229    async fn test_concurrent_misses_run_one_loader() {
230        const CONCURRENCY: usize = 8;
231
232        let hits = CACHE_HIT.with_label_values(&["pipeline_original"]);
233        let misses = CACHE_MISS.with_label_values(&["pipeline_original"]);
234        let before = (hits.get(), misses.get());
235        let cache = PipelineCache::new(Duration::from_secs(60));
236        let loads = AtomicUsize::new(0);
237        let release = Notify::new();
238        let mut requests = (0..CONCURRENCY)
239            .map(|_| {
240                Box::pin(cache.get_pipeline_str_with("db", "p", None, async {
241                    loads.fetch_add(1, Ordering::SeqCst);
242                    release.notified().await;
243                    Ok(content_at(1))
244                }))
245            })
246            .collect::<Vec<_>>();
247        for request in &mut requests {
248            assert!(futures::poll!(request.as_mut()).is_pending());
249        }
250        assert_eq!(loads.load(Ordering::SeqCst), 1);
251        release.notify_waiters();
252        for request in requests {
253            assert_eq!(request.await.unwrap(), content_at(1));
254        }
255        assert_eq!(loads.load(Ordering::SeqCst), 1);
256        assert_eq!(
257            (hits.get(), misses.get()),
258            (before.0 + CONCURRENCY as u64 - 1, before.1 + 1)
259        );
260    }
261
262    #[tokio::test]
263    async fn test_pipeline_lookup_metrics() {
264        let cache = PipelineCache::new(Duration::from_secs(60));
265        for compiled in [false, true] {
266            let cache_type = if compiled {
267                "pipeline"
268            } else {
269                "pipeline_original"
270            };
271            let hits = CACHE_HIT.with_label_values(&[cache_type]);
272            let misses = CACHE_MISS.with_label_values(&[cache_type]);
273            let before = (hits.get(), misses.get());
274            let load = async |fail| {
275                if compiled {
276                    cache
277                        .get_pipeline_with("db", "p", None, async {
278                            if fail {
279                                return Err(crate::error::PipelineNotFoundSnafu {
280                                    name: "p",
281                                    version: None,
282                                }
283                                .build());
284                            }
285                            Ok(Arc::new(
286                                crate::manager::table::PipelineTable::compile_pipeline(
287                                    "transform:\n  - fields: [message]\n    type: string",
288                                )?,
289                            ))
290                        })
291                        .await
292                        .map(|_| ())
293                } else {
294                    cache
295                        .get_pipeline_str_with("db", "p", None, async {
296                            if fail {
297                                return Err(crate::error::PipelineNotFoundSnafu {
298                                    name: "p",
299                                    version: None,
300                                }
301                                .build());
302                            }
303                            Ok(content_at(1))
304                        })
305                        .await
306                        .map(|_| ())
307                }
308            };
309            assert!(load(true).await.is_err());
310            load(false).await.unwrap();
311            // A warm lookup must not evaluate the failing initializer.
312            load(true).await.unwrap();
313            assert_eq!((hits.get(), misses.get()), (before.0 + 1, before.1 + 2));
314        }
315    }
316
317    #[tokio::test]
318    async fn test_failover_lookup_metrics() {
319        let hits = CACHE_HIT.with_label_values(&["pipeline_failover"]);
320        let misses = CACHE_MISS.with_label_values(&["pipeline_failover"]);
321        let before = (hits.get(), misses.get());
322        let cache = PipelineCache::new(Duration::from_secs(60));
323        assert!(
324            cache
325                .get_failover_cache("db", "p", None)
326                .await
327                .unwrap()
328                .is_none()
329        );
330        let local = PipelineContent {
331            schema: "a".into(),
332            ..content_at(1)
333        };
334        cache.insert_failover_cache(local.clone(), true).await;
335        // Direct schema lookup and cross-schema fallback each count one hit.
336        assert_eq!(
337            cache.get_failover_cache("a", "p", None).await.unwrap(),
338            Some(local.clone())
339        );
340        assert_eq!(
341            cache.get_failover_cache("db", "p", None).await.unwrap(),
342            Some(local)
343        );
344        cache
345            .insert_failover_cache(
346                PipelineContent {
347                    schema: "b".into(),
348                    ..content_at(2)
349                },
350                true,
351            )
352            .await;
353        assert!(cache.get_failover_cache("db", "p", None).await.is_err());
354        // Global pipelines win over local entries inserted before or after them.
355        for (schema, version) in [(EMPTY_SCHEMA_NAME, 3), ("c", 4)] {
356            cache
357                .insert_failover_cache(
358                    PipelineContent {
359                        schema: schema.into(),
360                        ..content_at(version)
361                    },
362                    true,
363                )
364                .await;
365            assert_eq!(
366                cache.get_failover_cache("db", "p", None).await.unwrap(),
367                Some(content_at(3))
368            );
369        }
370        assert_eq!((hits.get(), misses.get()), (before.0 + 4, before.1 + 2));
371    }
372
373    #[tokio::test]
374    async fn test_delete_drops_version_pinned_entry() {
375        let cache = PipelineCache::new(Duration::from_secs(60));
376        let content = content_at(1);
377        let version = Some(content.version);
378
379        cache
380            .get_pipeline_str_with("db", "p", version, async { Ok(content.clone()) })
381            .await
382            .unwrap();
383
384        cache.invalidate("p", version).await;
385
386        let loads = AtomicUsize::new(0);
387        cache
388            .get_pipeline_str_with("db", "p", version, async {
389                loads.fetch_add(1, Ordering::SeqCst);
390                Ok(content.clone())
391            })
392            .await
393            .unwrap();
394        assert_eq!(loads.load(Ordering::SeqCst), 1);
395    }
396
397    #[tokio::test]
398    async fn test_create_drops_stale_latest_and_primes_failover() {
399        let cache = PipelineCache::new(Duration::from_secs(60));
400        let v2 = content_at(2);
401
402        cache
403            .get_pipeline_str_with("a", "p", None, async { Ok(content_at(1)) })
404            .await
405            .unwrap();
406        cache.insert_failover_cache(content_at(1), true).await;
407
408        cache.on_pipeline_created(v2.clone()).await;
409
410        let loads = AtomicUsize::new(0);
411        let cached = cache
412            .get_pipeline_str_with("a", "p", None, async {
413                loads.fetch_add(1, Ordering::SeqCst);
414                Ok(v2.clone())
415            })
416            .await
417            .unwrap();
418        assert_eq!(loads.load(Ordering::SeqCst), 1);
419        assert_eq!(cached.version, v2.version);
420
421        let failover = cache.get_failover_cache("b", "p", None).await.unwrap();
422        assert_eq!(failover.map(|c| c.version), Some(v2.version));
423    }
424}
pipeline/manager/pipeline_cache.rs

pipeline/manager/
pipeline_cache.rs
