feat: allow spreading load across runners

Author: ctron

TODO.md

@@ -0,0 +1,4 @@
+# ToDo
+
+* [ ] Allowing skipping data part of the migration
+* [x] Allow concurrent instances (x of y)

docs/book/modules/migration/pages/data.adoc

@@ -0,0 +1,54 @@
+= Data migration guide
+
+In some cases, it is necessary to also migrate data during an upgrade. This may require the re-ingestion of all stored
+documents.
+
+This is the case, for example, a new field is added, which requires the extraction of information from the original
+document.
+
+== Strategy
+
+The overall strategy for this is:
+
+* Prevent access to the in-migration database
+* Create new database features (columns, …) in a compatible way (e.g. "allow nulls")
+* Process re-ingestion of documents
+* Modify database features to the target definition (e.g. "remove nullable")
+* Switch to new software version
+* Grant access to the new database structures
+
+This can be achieved with a read-only replica:
+
+* Prevent access to the in-migration database
+    * Create a read-only replica of the database
+    * Reconfigure trustify to use the read-only replica
+    * All mutable operations will fail, reporting `503 Service Unavailable`
+* Create new database features
+* Process re-ingestion of documents
+* Modify database features to the target definition
+* Switch to new software version
+* Grant access to the new database structures
+    * Switch to new database version
+    * Drop old, read-only replica
+
+== Running the re-ingestion
+
+The re-ingestion can be run in two ways:
+
+* During the SeoORM migration
+* Before the SeoORM migration
+
+The main difference is that when running during the SeoORM migration, you have less control over the process. It will
+be driven by the SeoORM migration, and you have to wait until everything is finished.
+
+Running before the SeoORM migration, you can, for example, run multiple instances of the re-ingestion. You can also run
+re-ingestion on a database copy, and then switch over to replace the database with the migrated version.
+
+Afterwards, you can run SeoORM migrations, which will then skip those DB modifications (as they are already applied) and
+also skip the re-ingestion.
+
+== The lazy way
+
+The lazy way, which is the default, will simply perform those steps during the SeaORM migration. However, there are a
+bunch of downsides. That's why it is not recommended for production setups. However, it may just work fine for small
+test setup. Making the process a lot easier.

migration/src/data/migration.rs

@@ -1,6 +1,6 @@
 use crate::{
     async_trait,
-    data::{Document, DocumentProcessor, Handler, Options},
+    data::{Document, DocumentProcessor, Handler, Options, Partition},
 };
 use clap::Parser;
 use futures::executor::block_on;
@@ -55,6 +55,7 @@ pub struct SchemaDataManager<'c> {
     pub manager: &'c SchemaManager<'c>,
     storage: &'c DispatchBackend,
     options: &'c Options,
+    partition: Partition,
 }
 
 impl<'c> SchemaDataManager<'c> {
@@ -63,10 +64,16 @@ impl<'c> SchemaDataManager<'c> {
         storage: &'c DispatchBackend,
         options: &'c Options,
     ) -> Self {
+        let partition = Partition {
+            current: options.current,
+            total: options.total,
+        };
+
         Self {
             manager,
             storage,
             options,
+            partition,
         }
     }
 

migration/src/data/mod.rs

@@ -1,7 +1,9 @@
 mod migration;
+mod partition;
 mod run;
 
 pub use migration::*;
+pub use partition::*;
 pub use run::*;
 
 use anyhow::{anyhow, bail};
@@ -14,7 +16,7 @@ use sea_orm::{
     ConnectionTrait, DatabaseTransaction, DbErr, EntityTrait, ModelTrait, TransactionTrait,
 };
 use sea_orm_migration::{MigrationTrait, SchemaManager};
-use std::num::NonZeroUsize;
+use std::num::{NonZeroU64, NonZeroUsize};
 use trustify_common::id::Id;
 use trustify_entity::{sbom, source_document};
 use trustify_module_storage::service::{StorageBackend, StorageKey, dispatch::DispatchBackend};
@@ -27,7 +29,7 @@ pub enum Sbom {
 
 #[allow(async_fn_in_trait)]
 pub trait Document: Sized + Send + Sync {
-    type Model: Send;
+    type Model: Partitionable + Send;
 
     async fn all<C>(tx: &C) -> Result<Vec<Self::Model>, DbErr>
     where
@@ -95,14 +97,22 @@ where
 
 #[derive(Clone, Debug, PartialEq, Eq, clap::Parser)]
 pub struct Options {
+    /// Number of concurrent documents being processes
     #[arg(long, env = "MIGRATION_DATA_CONCURRENT", default_value = "5")]
     pub concurrent: NonZeroUsize,
+
+    #[arg(long, env = "MIGRATION_DATA_CURRENT_RUNNER", default_value = "0")]
+    pub current: u64,
+    #[arg(long, env = "MIGRATION_DATA_TOTAL_RUNNER", default_value = "1")]
+    pub total: NonZeroU64,
 }
 
 impl Default for Options {
     fn default() -> Self {
         Self {
             concurrent: unsafe { NonZeroUsize::new_unchecked(5) },
+            current: 0,
+            total: unsafe { NonZeroU64::new_unchecked(1) },
         }
     }
 }
@@ -128,30 +138,34 @@ impl<'c> DocumentProcessor for SchemaManager<'c> {
     where
         D: Document,
     {
+        let partition = Partition::default();
         let db = self.get_connection();
 
         let tx = db.begin().await?;
         let all = D::all(&tx).await?;
         drop(tx);
 
-        stream::iter(all)
-            .map(async |model| {
-                let tx = db.begin().await?;
-
-                let doc = D::source(&model, storage, &tx).await.map_err(|err| {
-                    DbErr::Migration(format!("Failed to load source document: {err}"))
-                })?;
-                f.call(doc, model, &tx).await.map_err(|err| {
-                    DbErr::Migration(format!("Failed to process document: {err}"))
-                })?;
-
-                tx.commit().await?;
-
-                Ok::<_, DbErr>(())
-            })
-            .buffer_unordered(options.concurrent.into())
-            .try_collect::<Vec<_>>()
-            .await?;
+        stream::iter(
+            all.into_iter()
+                .filter(|model| partition.is_selected::<D>(&model)),
+        )
+        .map(async |model| {
+            let tx = db.begin().await?;
+
+            let doc = D::source(&model, storage, &tx).await.map_err(|err| {
+                DbErr::Migration(format!("Failed to load source document: {err}"))
+            })?;
+            f.call(doc, model, &tx)
+                .await
+                .map_err(|err| DbErr::Migration(format!("Failed to process document: {err}")))?;
+
+            tx.commit().await?;
+
+            Ok::<_, DbErr>(())
+        })
+        .buffer_unordered(options.concurrent.into())
+        .try_collect::<Vec<_>>()
+        .await?;
 
         Ok(())
     }

migration/src/data/partition.rs

@@ -0,0 +1,45 @@
+use crate::data::Document;
+use std::hash::{DefaultHasher, Hash, Hasher};
+use std::num::{NonZeroU64, NonZeroUsize};
+use trustify_entity::sbom;
+
+#[derive(Debug, Copy, Clone)]
+pub struct Partition {
+    pub current: u64,
+    pub total: NonZeroU64,
+}
+
+pub trait Partitionable {
+    fn hashed_id(&self) -> u64;
+}
+
+impl Partitionable for sbom::Model {
+    fn hashed_id(&self) -> u64 {
+        let mut hasher = DefaultHasher::new();
+        self.sbom_id.hash(&mut hasher);
+        hasher.finish()
+    }
+}
+
+impl Default for Partition {
+    fn default() -> Self {
+        Self::new_one()
+    }
+}
+
+impl Partition {
+    pub const fn new_one() -> Self {
+        Self {
+            current: 0,
+            total: unsafe { NonZeroU64::new_unchecked(1) },
+        }
+    }
+
+    pub fn is_selected<D>(&self, document: &D::Model) -> bool
+    where
+        D: Document,
+        D::Model: Partitionable,
+    {
+        document.hashed_id() % self.total == self.current
+    }
+}