Implement sharded mode execution

Summary:
## This stack

Running the worker for all repos is causing OOMs; sharding should help.

## This diff

With the preparation we've done eariler in the stack, we can now enabled sharded execution.

Reviewed By: singhsrb

Differential Revision: D64402510

fbshipit-source-id: 9b0a57a0423d47a5c48233fea63bd50e0cc5af1e

eden/mononoke/async_requests/async_requests_client/src/lib.rs

@@ -19,7 +19,7 @@ use mononoke_api::MononokeRepo;
 use mononoke_api::RepositoryId;
 use mononoke_app::MononokeApp;
 use requests_table::SqlLongRunningRequestsQueue;
-use slog::info;
+use slog::debug;
 use sql_construct::SqlConstructFromDatabaseConfig;
 use sql_ext::facebook::MysqlOptions;
 
@@ -46,7 +46,7 @@ pub async fn open_sql_connection(
 ) -> Result<SqlLongRunningRequestsQueue, Error> {
     let config = app.repo_configs().common.async_requests_config.clone();
     if let Some(config) = config.db_config {
-        info!(
+        debug!(
             app.logger(),
             "Initializing async_requests with an explicit config"
         );

eden/mononoke/async_requests/worker/Cargo.toml

@@ -14,6 +14,7 @@ async-trait = "0.1.71"
 async_requests = { version = "0.1.0", path = "../lib" }
 async_requests_client = { version = "0.1.0", path = "../async_requests_client" }
 async_requests_types_thrift = { version = "0.1.0", path = "../if" }
+blobstore = { version = "0.1.0", path = "../../blobstore" }
 clap = { version = "4.5.20", features = ["derive", "env", "string", "unicode", "wrap_help"] }
 cloned = { version = "0.1.0", git = "https://github.com/facebookexperimental/rust-shed.git", branch = "main" }
 cmdlib_logging = { version = "0.1.0", path = "../../cmdlib/log" }
@@ -31,6 +32,8 @@ mononoke_api = { version = "0.1.0", path = "../../mononoke_api" }
 mononoke_app = { version = "0.1.0", path = "../../cmdlib/mononoke_app" }
 mononoke_types = { version = "0.1.0", path = "../../mononoke_types" }
 repo_authorization = { version = "0.1.0", path = "../../repo_authorization" }
+requests_table = { version = "0.1.0", path = "../requests_table" }
+sharding_ext = { version = "0.1.0", path = "../../cmdlib/sharding_ext" }
 slog = { version = "2.7", features = ["max_level_trace", "nested-values"] }
 source_control = { version = "0.1.0", path = "../../scs/if" }
 stats = { version = "0.1.0", git = "https://github.com/facebookexperimental/rust-shed.git", branch = "main" }
@@ -39,4 +42,3 @@ tokio = { version = "1.37.0", features = ["full", "test-util", "tracing"] }
 [dev-dependencies]
 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" }
-requests_table = { version = "0.1.0", path = "../requests_table" }

eden/mononoke/async_requests/worker/TARGETS

@@ -25,12 +25,15 @@ rust_binary(
         "//common/rust/shed/stats:stats",
         "//eden/mononoke/async_requests:async_requests",
         "//eden/mononoke/async_requests:async_requests_client",
+        "//eden/mononoke/async_requests:requests_table",
         "//eden/mononoke/async_requests/if:async_requests_types-thrift-rust",
+        "//eden/mononoke/blobstore:blobstore",
         "//eden/mononoke/blobstore:ephemeral_blobstore",
         "//eden/mononoke/cmdlib:cmdlib_logging",
         "//eden/mononoke/cmdlib:environment",
         "//eden/mononoke/cmdlib/mononoke_app:mononoke_app",
         "//eden/mononoke/cmdlib/sharding:executor_lib",
+        "//eden/mononoke/cmdlib/sharding_ext:sharding_ext",
         "//eden/mononoke/megarepo_api:megarepo_api",
         "//eden/mononoke/megarepo_api:megarepo_error",
         "//eden/mononoke/metaconfig:metaconfig_types",

eden/mononoke/async_requests/worker/src/main.rs

@@ -14,17 +14,27 @@ use std::sync::atomic::AtomicBool;
 use std::sync::atomic::Ordering;
 use std::sync::Arc;
 
+use anyhow::Context;
 use anyhow::Result;
+use async_requests::AsyncMethodRequestQueue;
+use async_requests_client::open_blobstore;
+use async_requests_client::open_sql_connection;
+use async_trait::async_trait;
 use clap::Parser;
+use cloned::cloned;
 use cmdlib_logging::ScribeLoggingArgs;
+use context::CoreContext;
 use context::SessionContainer;
 use environment::BookmarkCacheDerivedData;
 use environment::BookmarkCacheKind;
 use environment::BookmarkCacheOptions;
+use executor_lib::args::ShardedExecutorArgs;
+use executor_lib::RepoShardedProcess;
 use executor_lib::RepoShardedProcessExecutor;
 use fbinit::FacebookInit;
 use megarepo_api::MegarepoApi;
 use metaconfig_types::ShardedService;
+use mononoke_api::Mononoke;
 use mononoke_api::Repo;
 use mononoke_app::args::HooksAppExtension;
 use mononoke_app::args::RepoFilterAppExtension;
@@ -33,12 +43,21 @@ use mononoke_app::args::WarmBookmarksCacheExtension;
 use mononoke_app::fb303::AliveService;
 use mononoke_app::fb303::Fb303AppExtension;
 use mononoke_app::MononokeAppBuilder;
+use mononoke_app::MononokeReposManager;
+use requests_table::SqlLongRunningRequestsQueue;
+use sharding_ext::RepoShard;
+use slog::info;
 
 const SERVICE_NAME: &str = "async_requests_worker";
 
+const SM_CLEANUP_TIMEOUT_SECS: u64 = 60;
+
 /// Processes the megarepo async requests
 #[derive(Parser)]
 struct AsyncRequestsWorkerArgs {
+    #[clap(flatten)]
+    pub sharded_executor_args: ShardedExecutorArgs,
+
     #[clap(flatten)]
     shutdown_timeout_args: ShutdownTimeoutArgs,
     #[clap(flatten)]
@@ -55,6 +74,75 @@ struct AsyncRequestsWorkerArgs {
     process_all_repos: bool,
 }
 
+pub struct WorkerProcess {
+    ctx: Arc<CoreContext>,
+    args: Arc<AsyncRequestsWorkerArgs>,
+    repos_mgr: Arc<MononokeReposManager<Repo>>,
+    mononoke: Arc<Mononoke<Repo>>,
+    megarepo: Arc<MegarepoApi<Repo>>,
+    sql_connection: Arc<SqlLongRunningRequestsQueue>,
+    blobstore: Arc<dyn blobstore::Blobstore>,
+    will_exit: Arc<AtomicBool>,
+}
+
+impl WorkerProcess {
+    pub(crate) fn new(
+        ctx: Arc<CoreContext>,
+        args: Arc<AsyncRequestsWorkerArgs>,
+        repos_mgr: Arc<MononokeReposManager<Repo>>,
+        mononoke: Arc<Mononoke<Repo>>,
+        megarepo: Arc<MegarepoApi<Repo>>,
+        sql_connection: Arc<SqlLongRunningRequestsQueue>,
+        blobstore: Arc<dyn blobstore::Blobstore>,
+        will_exit: Arc<AtomicBool>,
+    ) -> Self {
+        Self {
+            ctx,
+            args,
+            repos_mgr,
+            mononoke,
+            megarepo,
+            sql_connection,
+            blobstore,
+            will_exit,
+        }
+    }
+}
+
+#[async_trait]
+impl RepoShardedProcess for WorkerProcess {
+    async fn setup(&self, repo: &RepoShard) -> Result<Arc<dyn RepoShardedProcessExecutor>> {
+        let repo_name = repo.repo_name.as_str();
+        let logger = self.repos_mgr.repo_logger(repo_name);
+        info!(&logger, "Setting up repo {}", repo_name);
+
+        let repo = self
+            .repos_mgr
+            .add_repo(repo_name)
+            .await
+            .with_context(|| format!("Failure in setting up repo {}", repo_name))?;
+        let repos = vec![repo.repo_identity.id()];
+        info!(&logger, "Completed setup for repos {:?}", repos);
+
+        let queue = Arc::new(AsyncMethodRequestQueue::new(
+            self.sql_connection.clone(),
+            self.blobstore.clone(),
+            Some(repos),
+        ));
+
+        let executor = worker::AsyncMethodRequestWorker::new(
+            self.args.clone(),
+            self.ctx.clone(),
+            queue,
+            self.mononoke.clone(),
+            self.megarepo.clone(),
+            self.will_exit.clone(),
+        )
+        .await?;
+        Ok(Arc::new(executor))
+    }
+}
+
 #[fbinit::main]
 fn main(fb: FacebookInit) -> Result<()> {
     let app = MononokeAppBuilder::new(fb)
@@ -69,50 +157,91 @@ fn main(fb: FacebookInit) -> Result<()> {
         .build::<AsyncRequestsWorkerArgs>()?;
 
     let args: Arc<AsyncRequestsWorkerArgs> = Arc::new(app.args()?);
-
     let env = app.environment();
+    let logger = app.logger().clone();
     let runtime = app.runtime().clone();
     let session = SessionContainer::new_with_defaults(env.fb);
-    let ctx = session.new_context(app.logger().clone(), env.scuba_sample_builder.clone());
-
-    let mononoke = Arc::new(
-        runtime
-            .block_on(app.open_managed_repos::<Repo>(Some(ShardedService::AsyncRequestsWorker)))?
-            .make_mononoke_api()?,
-    );
-    let repos = mononoke.known_repo_ids();
+    let ctx = Arc::new(session.new_context(app.logger().clone(), env.scuba_sample_builder.clone()));
+
+    let service_name = Some(ShardedService::AsyncRequestsWorker);
+    let repos_mgr = Arc::new(runtime.block_on(app.open_managed_repos(service_name))?);
+    let mononoke = Arc::new(repos_mgr.make_mononoke_api()?);
     let megarepo = Arc::new(MegarepoApi::new(&app, mononoke.clone())?);
 
+    let sql_connection = Arc::new(runtime.block_on(open_sql_connection(fb, &app))?);
+    let blobstore = runtime.block_on(open_blobstore(fb, &app))?;
     let will_exit = Arc::new(AtomicBool::new(false));
-    let filter_repos = if args.process_all_repos {
-        None
-    } else {
-        Some(repos)
-    };
-    let queue = Arc::new(runtime.block_on(async_requests_client::build(fb, &app, filter_repos))?);
-    let worker = runtime.block_on(worker::AsyncMethodRequestWorker::new(
-        args.clone(),
-        Arc::new(ctx),
-        queue,
-        mononoke,
-        megarepo,
-        will_exit.clone(),
-    ))?;
 
     app.start_monitoring(SERVICE_NAME, AliveService)?;
     app.start_stats_aggregation()?;
 
-    let run_worker = { move |_app| async move { worker.execute().await } };
-
-    app.run_until_terminated(
-        run_worker,
-        move || will_exit.store(true, Ordering::Relaxed),
-        args.shutdown_timeout_args.shutdown_grace_period,
-        async {
-            // the code to gracefully stop things goes here
+    if let Some(mut executor) = args.sharded_executor_args.clone().build_executor(
+        app.fb,
+        runtime.clone(),
+        &logger,
+        || {
+            Arc::new(WorkerProcess::new(
+                ctx.clone(),
+                args.clone(),
+                repos_mgr.clone(),
+                mononoke.clone(),
+                megarepo.clone(),
+                sql_connection.clone(),
+                blobstore.clone(),
+                will_exit.clone(),
+            ))
         },
-        args.shutdown_timeout_args.shutdown_timeout,
-    )?;
+        true, // enable shard (repo) level healing
+        SM_CLEANUP_TIMEOUT_SECS,
+    )? {
+        info!(logger, "Starting sharded process");
+        // The Sharded Process Executor needs to branch off and execute
+        // on its own dedicated task spawned off the common tokio runtime.
+        runtime.spawn({
+            let logger = logger.clone();
+            {
+                cloned!(will_exit);
+                async move { executor.block_and_execute(&logger, will_exit).await }
+            }
+        });
+
+        app.wait_until_terminated(
+            move || will_exit.store(true, Ordering::Relaxed),
+            args.shutdown_timeout_args.shutdown_grace_period,
+            async {
+                info!(logger, "Shutdown");
+            },
+            args.shutdown_timeout_args.shutdown_timeout,
+        )?;
+    } else {
+        let logger = logger.clone();
+        let queue = Arc::new(AsyncMethodRequestQueue::new(
+            sql_connection,
+            blobstore,
+            None,
+        ));
+
+        info!(logger, "Starting unsharded executor for all repos");
+        let executor = runtime.block_on(worker::AsyncMethodRequestWorker::new(
+            args.clone(),
+            ctx.clone(),
+            queue.clone(),
+            mononoke.clone(),
+            megarepo.clone(),
+            will_exit.clone(),
+        ))?;
+        let run_worker = { move |_app| async move { executor.execute().await } };
+
+        app.run_until_terminated(
+            run_worker,
+            move || will_exit.store(true, Ordering::Relaxed),
+            args.shutdown_timeout_args.shutdown_grace_period,
+            async {
+                info!(logger, "Shutdown");
+            },
+            args.shutdown_timeout_args.shutdown_timeout,
+        )?;
+    }
 
     Ok(())
 }