Implement cached git ref content mapping using pub sub model

Summary:
With the work done so far, I can finally implement what I actually wanted :)

`git_ref_content_mapping` stores the mapping of refs (or bookmarks as they are known in Mononoke) to git shas of blobs or trees. Since refs pointings to content types is not standard in Mononoke, we have created this alternate route that bypasses the whole bookmark mechanism.

As part of Git read protocol, we end up fetching all the refs all the time unfortunately. Since `git_ref_content_mapping` does not have a cached implementation, this means queries to the DB. We can't directly cache it like other mappings (bonsai_git_mapping, bonsai_hg_mapping, etc.) because this mapping also supports updates. We need a WBC kind-of mechanism that keeps the cache up to date without incessantly polling the DB. `CachedGitRefContentMapping` using pub-sub scribed based notification mechanism is the answer.

Differential Revision: D75061565

fbshipit-source-id: 7dbabe0948d4bc9f7b6343a67fd7fb342a3f4cd7

Author: RajivTS

eden/mononoke/git_ref_content_mapping/BUCK

@@ -9,12 +9,20 @@ rust_library(
         "src/**/*.rs",
         "schemas/**/*.sql",
     ]),
+    named_deps = {"slog": "//common/rust/shed/tracing_slog_compat:tracing_slog_compat"},
     deps = [
         "fbsource//third-party/rust:anyhow",
+        "fbsource//third-party/rust:arc-swap",
         "fbsource//third-party/rust:async-trait",
+        "fbsource//third-party/rust:tokio",
+        "//common/rust/shed/cloned:cloned",
         "//common/rust/shed/facet:facet",
+        "//common/rust/shed/justknobs_stub:justknobs",
+        "//common/rust/shed/stats:stats",
         "//eden/mononoke/common/rust/sql_ext:sql_ext",
         "//eden/mononoke/common/sql_construct:sql_construct",
+        "//eden/mononoke/features/repo_update_logger:repo_update_logger",
+        "//eden/mononoke/mononoke_macros:mononoke_macros",
         "//eden/mononoke/mononoke_types:mononoke_types",
     ],
 )

eden/mononoke/git_ref_content_mapping/Cargo.toml

@@ -13,14 +13,21 @@ path = "test/main.rs"
 
 [dependencies]
 anyhow = "1.0.95"
+arc-swap = { version = "1.5", features = ["weak"] }
 async-trait = "0.1.86"
+cloned = { version = "0.1.0", git = "https://github.com/facebookexperimental/rust-shed.git", branch = "main" }
 facet = { version = "0.1.0", git = "https://github.com/facebookexperimental/rust-shed.git", branch = "main" }
+justknobs = { version = "0.1.0", git = "https://github.com/facebookexperimental/rust-shed.git", branch = "main" }
+mononoke_macros = { version = "0.1.0", path = "../mononoke_macros" }
 mononoke_types = { version = "0.1.0", path = "../mononoke_types" }
+repo_update_logger = { version = "0.1.0", path = "../features/repo_update_logger" }
+slog = { package = "tracing_slog_compat", version = "0.1.0", git = "https://github.com/facebookexperimental/rust-shed.git", branch = "main" }
 sql_construct = { version = "0.1.0", path = "../common/sql_construct" }
 sql_ext = { version = "0.1.0", path = "../common/rust/sql_ext" }
+stats = { version = "0.1.0", git = "https://github.com/facebookexperimental/rust-shed.git", branch = "main" }
+tokio = { version = "1.45.0", features = ["full", "test-util", "tracing"] }
 
 [dev-dependencies]
 fbinit = { version = "0.2.0", git = "https://github.com/facebookexperimental/rust-shed.git", branch = "main" }
 fbinit-tokio = { version = "0.1.2", git = "https://github.com/facebookexperimental/rust-shed.git", branch = "main" }
-mononoke_macros = { version = "0.1.0", path = "../mononoke_macros" }
 mononoke_types-mocks = { version = "0.1.0", path = "../mononoke_types/mocks" }

eden/mononoke/git_ref_content_mapping/src/cache.rs

@@ -0,0 +1,173 @@
+/*
+ * Copyright (c) Meta Platforms, Inc. and affiliates.
+ *
+ * This software may be used and distributed according to the terms of the
+ * GNU General Public License version 2.
+ */
+
+use std::sync::Arc;
+
+use anyhow::Context;
+use anyhow::Result;
+use arc_swap::ArcSwap;
+use async_trait::async_trait;
+use cloned::cloned;
+use mononoke_macros::mononoke;
+use mononoke_types::RepositoryId;
+use repo_update_logger::GitContentRefInfo;
+use slog::error;
+use slog::info;
+use stats::define_stats;
+use stats::prelude::TimeseriesStatic;
+use tokio::sync::broadcast::Receiver;
+use tokio::task::JoinHandle;
+
+use crate::GitRefContentMapping;
+use crate::GitRefContentMappingEntry;
+
+type Swappable<T> = Arc<ArcSwap<T>>;
+
+define_stats! {
+    prefix = "mononoke.content_refs_cache_update";
+    update_failure_count: timeseries(Average, Sum, Count),
+}
+
+pub struct CachedGitRefContentMapping {
+    inner: Arc<dyn GitRefContentMapping>,
+    entries: Swappable<Vec<GitRefContentMappingEntry>>,
+    updater_task: JoinHandle<()>,
+}
+
+#[allow(dead_code)]
+impl CachedGitRefContentMapping {
+    pub async fn new(
+        inner: Arc<dyn GitRefContentMapping>,
+        update_notification_receiver: Receiver<GitContentRefInfo>,
+        logger: slog::Logger,
+    ) -> Result<Self> {
+        let initial_entries = inner
+            .get_all_entries()
+            .await
+            .context("Error while getting initial set of git ref content mapping entries")?;
+        let entries = Arc::new(ArcSwap::from_pointee(initial_entries));
+        let updater_task = mononoke::spawn_task({
+            cloned!(entries, inner);
+            update_cache(entries, inner, update_notification_receiver, logger)
+        });
+        Ok(Self {
+            inner,
+            entries,
+            updater_task,
+        })
+    }
+}
+
+async fn update_cache(
+    entries: Swappable<Vec<GitRefContentMappingEntry>>,
+    bonsai_tag_mapping: Arc<dyn GitRefContentMapping>,
+    mut update_notification_receiver: Receiver<GitContentRefInfo>,
+    logger: slog::Logger,
+) {
+    loop {
+        let fallible_notification = update_notification_receiver
+            .recv()
+            .await
+            .context("Error while receiving update notification");
+        match fallible_notification {
+            Ok(_) => {
+                info!(
+                    logger,
+                    "Received update notification from scribe for updating content refs cache"
+                );
+                match bonsai_tag_mapping.get_all_entries().await {
+                    Ok(new_entries) => {
+                        let new_entries = Arc::new(new_entries);
+                        entries.store(new_entries);
+                        info!(
+                            logger,
+                            "Successfully updated the cache with new entries from the inner git ref content mapping"
+                        );
+                    }
+                    Err(e) => {
+                        error!(
+                            logger,
+                            "Failure in updating the cache with new entries from the inner git ref content mapping: {:?}",
+                            e
+                        );
+                        STATS::update_failure_count.add_value(1);
+                    }
+                }
+            }
+            Err(e) => {
+                error!(
+                    logger,
+                    "Failure in receiving notification from tags scribe category. Error: {:?}", e
+                );
+                STATS::update_failure_count.add_value(1);
+            }
+        }
+    }
+}
+
+impl Drop for CachedGitRefContentMapping {
+    fn drop(&mut self) {
+        // Need to abort the task before dropping the cache mapping
+        self.updater_task.abort();
+    }
+}
+
+#[async_trait]
+impl GitRefContentMapping for CachedGitRefContentMapping {
+    /// The repository for which this mapping has been created
+    fn repo_id(&self) -> RepositoryId {
+        self.inner.repo_id()
+    }
+
+    /// Fetch all the tag mapping entries for the given repo
+    async fn get_all_entries(&self) -> Result<Vec<GitRefContentMappingEntry>> {
+        if justknobs::eval(
+            "scm/mononoke:enable_git_ref_content_mapping_caching",
+            None,
+            None,
+        )
+        .unwrap_or(false)
+        {
+            Ok(self.entries.load_full().to_vec())
+        } else {
+            self.inner.get_all_entries().await
+        }
+    }
+
+    async fn get_entry_by_ref_name(
+        &self,
+        ref_name: String,
+    ) -> Result<Option<GitRefContentMappingEntry>> {
+        if justknobs::eval(
+            "scm/mononoke:enable_git_ref_content_mapping_caching",
+            None,
+            None,
+        )
+        .unwrap_or(false)
+        {
+            let entry = self
+                .entries
+                .load()
+                .iter()
+                .find(|&entry| entry.ref_name == ref_name)
+                .cloned();
+            Ok(entry)
+        } else {
+            self.inner.get_entry_by_ref_name(ref_name).await
+        }
+    }
+
+    async fn add_or_update_mappings(&self, entries: Vec<GitRefContentMappingEntry>) -> Result<()> {
+        // Writes are directly delegated to inner git ref content mapping
+        self.inner.add_or_update_mappings(entries).await
+    }
+
+    async fn delete_mappings_by_name(&self, ref_names: Vec<String>) -> Result<()> {
+        // Writes are directly delegated to inner git ref content mapping
+        self.inner.delete_mappings_by_name(ref_names).await
+    }
+}

eden/mononoke/git_ref_content_mapping/src/lib.rs

@@ -5,13 +5,15 @@
  * GNU General Public License version 2.
  */
 
+mod cache;
 mod sql;
 
 use anyhow::Result;
 use async_trait::async_trait;
 use mononoke_types::RepositoryId;
 use mononoke_types::hash::GitSha1;
 
+pub use crate::cache::CachedGitRefContentMapping;
 pub use crate::sql::SqlGitRefContentMapping;
 pub use crate::sql::SqlGitRefContentMappingBuilder;
 

eden/mononoke/repo_factory/src/lib.rs

@@ -114,6 +114,7 @@ use filestore::ArcFilestoreConfig;
 use filestore::FilestoreConfig;
 use futures_watchdog::WatchdogExt;
 use git_ref_content_mapping::ArcGitRefContentMapping;
+use git_ref_content_mapping::CachedGitRefContentMapping;
 use git_ref_content_mapping::SqlGitRefContentMappingBuilder;
 use git_source_of_truth::ArcGitSourceOfTruthConfig;
 use git_source_of_truth::SqlGitSourceOfTruthConfigBuilder;
@@ -1122,14 +1123,39 @@ impl RepoFactory {
         &self,
         repo_config: &ArcRepoConfig,
         repo_identity: &ArcRepoIdentity,
+        repo_event_publisher: &ArcRepoEventPublisher,
     ) -> Result<ArcGitRefContentMapping> {
         let git_ref_content_mapping = self
             .open_sql::<SqlGitRefContentMappingBuilder>(repo_config)
             .await
             .context(RepoFactoryError::GitRefContentMapping)?
             .build(repo_identity.id());
-        // TODO(rajshar): Add caching for git_ref_content_mapping
-        Ok(Arc::new(git_ref_content_mapping))
+        let repo_name = repo_identity.name();
+        if justknobs::eval(
+            "scm/mononoke:enable_git_ref_content_mapping_caching",
+            None,
+            Some(repo_name),
+        )
+        .unwrap_or(false)
+        {
+            let logger = self.env.logger.clone();
+            match repo_event_publisher.subscribe_for_content_refs_updates(&repo_name.to_string()) {
+                Ok(update_notification_receiver) => {
+                    let cached_git_ref_content_mapping = CachedGitRefContentMapping::new(
+                        Arc::new(git_ref_content_mapping),
+                        update_notification_receiver,
+                        logger,
+                    )
+                    .await?;
+                    Ok(Arc::new(cached_git_ref_content_mapping))
+                }
+                // The scribe configuration does not exist for content ref updates for this repo, so use the non-cached
+                // version of git_ref_content_mapping
+                Err(_) => Ok(Arc::new(git_ref_content_mapping)),
+            }
+        } else {
+            Ok(Arc::new(git_ref_content_mapping))
+        }
     }
 
     pub async fn git_bundle_uri(

eden/mononoke/tests/integration/mononoke_git_server/test-mononoke-git-server-push-with-tree-and-blob-ref.t

@@ -5,13 +5,22 @@
 # directory of this source tree.
 
   $ export GIT_CONTENT_REFS_SCRIBE_CATEGORY=mononoke_git_content_refs
+  $ export CONTENT_REFS_SCRIBE_CATEGORY=mononoke_git_content_refs
   $ export MONONOKE_TEST_SCRIBE_LOGGING_DIRECTORY=$TESTTMP/scribe_logs/
   $ . "${TEST_FIXTURES}/library.sh"
   $ REPOTYPE="blob_files"
   $ setup_common_config $REPOTYPE
   $ GIT_REPO_ORIGIN="${TESTTMP}/origin/repo-git"
   $ GIT_REPO="${TESTTMP}/repo-git"
 
+  $ merge_just_knobs <<EOF
+  > {
+  >   "bools": {
+  >     "scm/mononoke:enable_git_ref_content_mapping_caching": true
+  >   }
+  > }
+  > EOF
+
 # Setup git repository
   $ mkdir -p "$GIT_REPO_ORIGIN"
   $ cd "$GIT_REPO_ORIGIN"