cache fakeip

Author: bdbai

ytflow/src/config/plugin/fakeip.rs

@@ -2,6 +2,7 @@ use serde::Deserialize;
 
 use crate::config::factory::*;
 use crate::config::*;
+use crate::data::PluginId;
 
 #[cfg_attr(not(feature = "plugins"), allow(dead_code))]
 #[derive(Clone, Deserialize)]
@@ -10,12 +11,17 @@ pub struct FakeIpFactory<'a> {
     prefix_v6: [u8; 14],
     // Reserved for CNAME, TXT and SRV support
     fallback: &'a str,
+    #[serde(skip)]
+    plugin_id: Option<PluginId>,
 }
 
 impl<'de> FakeIpFactory<'de> {
     pub(in super::super) fn parse(plugin: &'de Plugin) -> ConfigResult<ParsedPlugin<'de, Self>> {
-        let Plugin { name, param, .. } = plugin;
-        let config: Self = parse_param(name, param)?;
+        let Plugin {
+            name, param, id, ..
+        } = plugin;
+        let mut config: Self = parse_param(name, param)?;
+        config.plugin_id = *id;
         Ok(ParsedPlugin {
             factory: config.clone(),
             requires: vec![Descriptor {
@@ -34,12 +40,27 @@ impl<'de> FakeIpFactory<'de> {
 impl<'de> Factory for FakeIpFactory<'de> {
     #[cfg(feature = "plugins")]
     fn load(&mut self, plugin_name: String, set: &mut PartialPluginSet) -> LoadResult<()> {
-        use crate::plugin::fakeip;
+        use crate::{data::PluginCache, plugin::fakeip};
 
-        let resolver = Arc::new(fakeip::FakeIp::new(self.prefix_v4, self.prefix_v6));
+        let db = set
+            .db
+            .ok_or_else(|| LoadError::DatabaseRequired {
+                plugin: plugin_name.clone(),
+            })?
+            .clone();
+        let cache = PluginCache::new(
+            self.plugin_id.ok_or_else(|| LoadError::DatabaseRequired {
+                plugin: plugin_name.clone(),
+            })?,
+            Some(db.clone()),
+        );
+        let plugin = Arc::new(fakeip::FakeIp::new(self.prefix_v4, self.prefix_v6, cache));
+        set.fully_constructed
+            .long_running_tasks
+            .push(tokio::spawn(fakeip::cache_writer(plugin.clone())));
         set.fully_constructed
             .resolver
-            .insert(plugin_name + ".resolver", resolver);
+            .insert(plugin_name + ".resolver", plugin);
         Ok(())
     }
 }

ytflow/src/plugin/fakeip.rs

@@ -1,48 +1,93 @@
+use std::collections::BTreeMap;
 use std::num::NonZeroUsize;
-use std::sync::Mutex;
+use std::sync::{Arc, Mutex};
 
 use async_trait::async_trait;
 use lru::LruCache;
+use serde::{Deserialize, Serialize};
 use smallvec::smallvec;
+use tokio::sync::Notify;
 
+use crate::data::PluginCache;
 use crate::flow::*;
 
 const CACHE_SIZE: NonZeroUsize = NonZeroUsize::new(1000).unwrap();
+const PLUGIN_CACHE_KEY: &str = "map";
 
 struct Inner {
     current: u16,
     cache: LruCache<String, u16>,
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct InnerCache {
+    current: u16,
+    cache: BTreeMap<String, u16>,
+}
+
 pub struct FakeIp {
     prefix_v4: u16,
     prefix_v6: [u8; 14],
-    inner: Mutex<Inner>,
+    inner: Arc<Mutex<Inner>>,
+    plugin_cache: PluginCache,
+    new_notify: Arc<Notify>,
 }
 
 impl FakeIp {
-    pub fn new(prefix_v4: [u8; 2], prefix_v6: [u8; 14]) -> Self {
-        // TODO: persist current cursor into cache
+    pub fn new(prefix_v4: [u8; 2], prefix_v6: [u8; 14], plugin_cache: PluginCache) -> Self {
+        let mut lru = LruCache::new(CACHE_SIZE);
+        let inner = match plugin_cache
+            .get::<InnerCache>(PLUGIN_CACHE_KEY)
+            .ok()
+            .flatten()
+        {
+            Some(cache) => {
+                for (k, v) in cache.cache {
+                    lru.put(k, v);
+                }
+                Inner {
+                    current: cache.current,
+                    cache: lru,
+                }
+            }
+            None => Inner {
+                current: 1,
+                cache: lru,
+            },
+        };
         Self {
             prefix_v4: u16::from_be_bytes(prefix_v4),
             prefix_v6,
-            inner: Mutex::new(Inner {
-                current: 1,
-                cache: LruCache::new(CACHE_SIZE),
-            }),
+            inner: Arc::new(Mutex::new(inner)),
+            plugin_cache,
+            new_notify: Arc::new(Notify::new()),
         }
     }
     fn lookup_or_alloc(&self, domain: String) -> u16 {
-        let mut inner = self.inner.lock().unwrap();
-        let cached = inner.cache.get(&*domain).copied();
-        if let Some(cached) = cached {
-            return cached;
-        }
-        let ret = inner.current;
-        inner.cache.put(domain, ret);
-        inner.current += 1;
+        let ret = {
+            let mut inner = self.inner.lock().unwrap();
+            let cached = inner.cache.get(&*domain).copied();
+            if let Some(cached) = cached {
+                return cached;
+            }
+            let ret = inner.current;
+            inner.cache.put(domain, ret);
+            inner.current = inner.current.wrapping_add(1);
+            ret
+        };
+        self.new_notify.notify_one();
         ret
     }
+    fn save_cache(&self) {
+        let cache = {
+            let inner = self.inner.lock().unwrap();
+            InnerCache {
+                current: inner.current,
+                cache: inner.cache.iter().map(|(k, v)| (k.clone(), *v)).collect(),
+            }
+        };
+        self.plugin_cache.set(PLUGIN_CACHE_KEY, &cache).ok();
+    }
 }
 
 #[async_trait]
@@ -62,3 +107,43 @@ impl Resolver for FakeIp {
         Ok(smallvec![bytes.into()])
     }
 }
+
+impl Drop for FakeIp {
+    fn drop(&mut self) {
+        self.save_cache();
+    }
+}
+
+pub async fn cache_writer(plugin: Arc<FakeIp>) {
+    let (plugin, notify) = {
+        let notify = plugin.new_notify.clone();
+        let weak = Arc::downgrade(&plugin);
+        drop(plugin);
+        (weak, notify)
+    };
+    if plugin.strong_count() == 0 {
+        panic!("fakeip has no strong reference left for cache_writer");
+    }
+
+    use tokio::select;
+    use tokio::time::{sleep, Duration};
+    loop {
+        let mut notified_fut = notify.notified();
+        let mut sleep_fut = sleep(Duration::from_secs(3600));
+        'debounce: loop {
+            select! {
+                _ = notified_fut => {
+                    notified_fut = notify.notified();
+                    sleep_fut = sleep(Duration::from_secs(3));
+                }
+                _ = sleep_fut => {
+                    break 'debounce;
+                }
+            }
+        }
+        match plugin.upgrade() {
+            Some(plugin) => plugin.save_cache(),
+            None => break,
+        }
+    }
+}