PostgresKeyValueStorage and NamespaceRoutedKeyValueStorage

golem-worker-executor/db/migration/postgres/keyvalue/postgres/001_init.sql

@@ -0,0 +1,25 @@
+CREATE TABLE kv_storage (
+    namespace TEXT NOT NULL,
+    key TEXT NOT NULL,
+    value BYTEA NOT NULL,
+    PRIMARY KEY (namespace, key)
+);
+
+CREATE TABLE set_storage (
+    namespace TEXT NOT NULL,
+    key TEXT NOT NULL,
+    value BYTEA NOT NULL,
+    PRIMARY KEY (namespace, key, value)
+);
+
+CREATE TABLE sorted_set_storage (
+    namespace TEXT NOT NULL,
+    key TEXT NOT NULL,
+    value_hash BYTEA NOT NULL,
+    value BYTEA NOT NULL,
+    score DOUBLE PRECISION NOT NULL,
+    PRIMARY KEY (namespace, key, value_hash)
+);
+
+CREATE INDEX idx_sorted_set_storage_namespace_key_score
+    ON sorted_set_storage (namespace, key, score);

golem-worker-executor/src/lib.rs

@@ -40,6 +40,7 @@ use crate::services::component::ComponentService;
 use crate::services::events::Events;
 use crate::services::golem_config::{
     EngineConfig, GolemConfig, HttpClientConfig, IndexedStorageConfig, KeyValueStorageConfig,
+    KeyValueStorageInnerConfig,
 };
 use crate::services::key_value::{DefaultKeyValueService, KeyValueService};
 use crate::services::oplog::plugin::{
@@ -71,6 +72,8 @@ use crate::storage::indexed::sqlite::SqliteIndexedStorage;
 use crate::storage::indexed::IndexedStorage;
 use crate::storage::keyvalue::memory::InMemoryKeyValueStorage;
 use crate::storage::keyvalue::multi_sqlite::MultiSqliteKeyValueStorage;
+use crate::storage::keyvalue::namespace_routed::NamespaceRoutedKeyValueStorage;
+use crate::storage::keyvalue::postgres::PostgresKeyValueStorage;
 use crate::storage::keyvalue::redis::RedisKeyValueStorage;
 use crate::storage::keyvalue::KeyValueStorage;
 use crate::workerctx::WorkerCtx;
@@ -357,6 +360,30 @@ pub async fn create_worker_executor_impl<Ctx: WorkerCtx, A: Bootstrap<Ctx> + ?Si
                 Arc::new(RedisKeyValueStorage::new(pool.clone()));
             (Some(pool), None, key_value_storage)
         }
+        KeyValueStorageConfig::Postgres(postgres) => {
+            let key_value_storage: Arc<dyn KeyValueStorage + Send + Sync> = Arc::new(
+                PostgresKeyValueStorage::configured(postgres)
+                    .await
+                    .map_err(|err| anyhow!(err))?,
+            );
+            (None, None, key_value_storage)
+        }
+        KeyValueStorageConfig::NamespaceRouted(namespace_routed) => {
+            let (cache_redis, cache_sqlite, cache_storage) =
+                build_inner_key_value_storage(&namespace_routed.cache).await?;
+            let (persistent_redis, persistent_sqlite, persistent_storage) =
+                build_inner_key_value_storage(&namespace_routed.persistent).await?;
+
+            let key_value_storage: Arc<dyn KeyValueStorage + Send + Sync> = Arc::new(
+                NamespaceRoutedKeyValueStorage::new(cache_storage, persistent_storage),
+            );
+
+            (
+                cache_redis.or(persistent_redis),
+                cache_sqlite.or(persistent_sqlite),
+                key_value_storage,
+            )
+        }
         KeyValueStorageConfig::InMemory(_) => {
             (None, None, Arc::new(InMemoryKeyValueStorage::new()))
         }
@@ -699,3 +726,58 @@ pub async fn create_worker_executor_impl<Ctx: WorkerCtx, A: Bootstrap<Ctx> + ?Si
 
     Ok((all, epoch_thread, epoch_stop))
 }
+
+async fn build_inner_key_value_storage(
+    config: &KeyValueStorageInnerConfig,
+) -> Result<
+    (
+        Option<RedisPool>,
+        Option<SqlitePool>,
+        Arc<dyn KeyValueStorage + Send + Sync>,
+    ),
+    anyhow::Error,
+> {
+    match config {
+        KeyValueStorageInnerConfig::Redis(redis) => {
+            let pool = RedisPool::configured(redis)
+                .await
+                .map_err(|err| anyhow!(err))?;
+            let key_value_storage: Arc<dyn KeyValueStorage + Send + Sync> =
+                Arc::new(RedisKeyValueStorage::new(pool.clone()));
+            Ok((Some(pool), None, key_value_storage))
+        }
+        KeyValueStorageInnerConfig::Postgres(postgres) => {
+            let key_value_storage: Arc<dyn KeyValueStorage + Send + Sync> = Arc::new(
+                PostgresKeyValueStorage::configured(postgres)
+                    .await
+                    .map_err(|err| anyhow!(err))?,
+            );
+            Ok((None, None, key_value_storage))
+        }
+        KeyValueStorageInnerConfig::InMemory(_) => {
+            let key_value_storage: Arc<dyn KeyValueStorage + Send + Sync> =
+                Arc::new(InMemoryKeyValueStorage::new());
+            Ok((None, None, key_value_storage))
+        }
+        KeyValueStorageInnerConfig::Sqlite(sqlite) => {
+            let pool = SqlitePool::configured(sqlite)
+                .await
+                .map_err(|err| anyhow!(err))?;
+            let key_value_storage: Arc<dyn KeyValueStorage + Send + Sync> = Arc::new(
+                SqliteKeyValueStorage::new(pool.clone())
+                    .await
+                    .map_err(|err| anyhow!(err))?,
+            );
+            Ok((None, Some(pool), key_value_storage))
+        }
+        KeyValueStorageInnerConfig::MultiSqlite(multi_sqlite) => {
+            let key_value_storage: Arc<dyn KeyValueStorage + Send + Sync> =
+                Arc::new(MultiSqliteKeyValueStorage::new(
+                    &multi_sqlite.root_dir,
+                    multi_sqlite.max_connections,
+                    multi_sqlite.foreign_keys,
+                ));
+            Ok((None, None, key_value_storage))
+        }
+    }
+}

golem-worker-executor/src/services/golem_config.rs

@@ -555,6 +555,8 @@ impl SafeDisplay for OplogConfig {
 #[serde(tag = "type", content = "config")]
 pub enum KeyValueStorageConfig {
     Redis(RedisConfig),
+    Postgres(KeyValueStoragePostgresConfig),
+    NamespaceRouted(KeyValueStorageNamespaceRoutedConfig),
     Sqlite(DbSqliteConfig),
     MultiSqlite(KeyValueStorageMultiSqliteConfig),
     InMemory(KeyValueStorageInMemoryConfig),
@@ -568,6 +570,14 @@ impl SafeDisplay for KeyValueStorageConfig {
                 let _ = writeln!(&mut result, "redis:");
                 let _ = writeln!(&mut result, "{}", inner.to_safe_string_indented());
             }
+            KeyValueStorageConfig::Postgres(inner) => {
+                let _ = writeln!(&mut result, "postgres:");
+                let _ = writeln!(&mut result, "{}", inner.to_safe_string_indented());
+            }
+            KeyValueStorageConfig::NamespaceRouted(inner) => {
+                let _ = writeln!(&mut result, "namespace-routed:");
+                let _ = writeln!(&mut result, "{}", inner.to_safe_string_indented());
+            }
             KeyValueStorageConfig::Sqlite(inner) => {
                 let _ = writeln!(&mut result, "sqlite:");
                 let _ = writeln!(&mut result, "{}", inner.to_safe_string_indented());
@@ -585,6 +595,74 @@ impl SafeDisplay for KeyValueStorageConfig {
     }
 }
 
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct KeyValueStoragePostgresConfig {
+    #[serde(flatten)]
+    pub postgres: DbPostgresConfig,
+}
+
+impl SafeDisplay for KeyValueStoragePostgresConfig {
+    fn to_safe_string(&self) -> String {
+        self.postgres.to_safe_string()
+    }
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+pub struct KeyValueStorageNamespaceRoutedConfig {
+    pub cache: KeyValueStorageInnerConfig,
+    pub persistent: KeyValueStorageInnerConfig,
+}
+
+impl SafeDisplay for KeyValueStorageNamespaceRoutedConfig {
+    fn to_safe_string(&self) -> String {
+        let mut result = String::new();
+        let _ = writeln!(&mut result, "cache:");
+        let _ = writeln!(&mut result, "{}", self.cache.to_safe_string_indented());
+        let _ = writeln!(&mut result, "persistent:");
+        let _ = writeln!(&mut result, "{}", self.persistent.to_safe_string_indented());
+        result
+    }
+}
+
+#[derive(Clone, Debug, Serialize, Deserialize)]
+#[serde(tag = "type", content = "config")]
+pub enum KeyValueStorageInnerConfig {
+    Redis(RedisConfig),
+    Postgres(KeyValueStoragePostgresConfig),
+    Sqlite(DbSqliteConfig),
+    MultiSqlite(KeyValueStorageMultiSqliteConfig),
+    InMemory(KeyValueStorageInMemoryConfig),
+}
+
+impl SafeDisplay for KeyValueStorageInnerConfig {
+    fn to_safe_string(&self) -> String {
+        let mut result = String::new();
+        match self {
+            KeyValueStorageInnerConfig::Redis(inner) => {
+                let _ = writeln!(&mut result, "redis:");
+                let _ = writeln!(&mut result, "{}", inner.to_safe_string_indented());
+            }
+            KeyValueStorageInnerConfig::Postgres(inner) => {
+                let _ = writeln!(&mut result, "postgres:");
+                let _ = writeln!(&mut result, "{}", inner.to_safe_string_indented());
+            }
+            KeyValueStorageInnerConfig::Sqlite(inner) => {
+                let _ = writeln!(&mut result, "sqlite:");
+                let _ = writeln!(&mut result, "{}", inner.to_safe_string_indented());
+            }
+            KeyValueStorageInnerConfig::MultiSqlite(inner) => {
+                let _ = writeln!(&mut result, "multi-sqlite:");
+                let _ = writeln!(&mut result, "{}", inner.to_safe_string_indented());
+            }
+            KeyValueStorageInnerConfig::InMemory(inner) => {
+                let _ = writeln!(&mut result, "in-memory:");
+                let _ = writeln!(&mut result, "{}", inner.to_safe_string_indented());
+            }
+        }
+        result
+    }
+}
+
 #[derive(Clone, Debug, Serialize, Deserialize)]
 pub struct KeyValueStorageMultiSqliteConfig {
     pub root_dir: PathBuf,

golem-worker-executor/src/storage/indexed/postgres.rs

@@ -24,7 +24,8 @@ use include_dir::include_dir;
 use sqlx::{Postgres, QueryBuilder};
 use std::time::Duration;
 
-static DB_MIGRATIONS: include_dir::Dir = include_dir!("$CARGO_MANIFEST_DIR/db/migration");
+static DB_MIGRATIONS: include_dir::Dir =
+    include_dir!("$CARGO_MANIFEST_DIR/db/migration/postgres/indexed");
 
 #[derive(Debug, Clone)]
 pub struct PostgresIndexedStorage {
@@ -33,6 +34,8 @@ pub struct PostgresIndexedStorage {
 }
 
 impl PostgresIndexedStorage {
+    const APPEND_MANY_CHUNK_SIZE: usize = 1024;
+
     pub async fn configured(config: &IndexedStoragePostgresConfig) -> Result<Self, String> {
         if config.drop_prefix_delete_batch_size == 0 {
             return Err(
@@ -239,20 +242,23 @@ impl IndexedStorage for PostgresIndexedStorage {
         self.pool
             .with_tx(svc_name, api_name, |tx| {
                 async move {
-                    let mut query_builder = QueryBuilder::<Postgres>::new(
-                        "INSERT INTO index_storage (namespace, key, id, value) ",
-                    );
-
-                    query_builder.push_values(converted_pairs, |mut builder, (id, value)| {
-                        builder
-                            .push_bind(namespace.clone())
-                            .push_bind(key.clone())
-                            .push_bind(id)
-                            .push_bind(value);
-                    });
-
-                    let query = query_builder.build();
-                    tx.execute(query).await?;
+                    for chunk in converted_pairs.chunks(Self::APPEND_MANY_CHUNK_SIZE) {
+                        let mut query_builder = QueryBuilder::<Postgres>::new(
+                            "INSERT INTO index_storage (namespace, key, id, value) ",
+                        );
+
+                        query_builder.push_values(chunk.iter(), |mut builder, (id, value)| {
+                            builder
+                                .push_bind(namespace.clone())
+                                .push_bind(key.clone())
+                                .push_bind(*id)
+                                .push_bind(value);
+                        });
+
+                        let query = query_builder.build();
+                        tx.execute(query).await?;
+                    }
+
                     Ok(())
                 }
                 .boxed()
@@ -413,27 +419,30 @@ impl IndexedStorage for PostgresIndexedStorage {
             self.drop_prefix_delete_batch_size,
             "drop_prefix_delete_batch_size",
         )?;
-        let mut deleted_rows = self.drop_prefix_delete_batch_size;
-
-        while deleted_rows >= self.drop_prefix_delete_batch_size {
-            let query = sqlx::query(
-                "WITH rows AS (SELECT ctid FROM index_storage WHERE namespace = $1 AND key = $2 AND id <= $3 ORDER BY id LIMIT $4) DELETE FROM index_storage t USING rows WHERE t.ctid = rows.ctid;",
-            )
-            .bind(&namespace)
-            .bind(&key)
-            .bind(last_dropped_id)
-            .bind(batch_size_i64);
-
-            deleted_rows = self
-                .pool
-                .with_rw(svc_name, api_name)
-                .execute(query)
-                .await
-                .map_err(|err| err.to_safe_string())?
-                .rows_affected();
-        }
+        let delete_batch_size = self.drop_prefix_delete_batch_size;
+        self.pool
+            .with_tx(svc_name, api_name, |tx| {
+                async move {
+                    let mut deleted_rows = delete_batch_size;
+
+                    while deleted_rows >= delete_batch_size {
+                        let query = sqlx::query(
+                            "WITH rows AS (SELECT ctid FROM index_storage WHERE namespace = $1 AND key = $2 AND id <= $3 ORDER BY id LIMIT $4) DELETE FROM index_storage t USING rows WHERE t.ctid = rows.ctid;",
+                        )
+                        .bind(&namespace)
+                        .bind(&key)
+                        .bind(last_dropped_id)
+                        .bind(batch_size_i64);
 
-        Ok(())
+                        deleted_rows = tx.execute(query).await?.rows_affected();
+                    }
+
+                    Ok(())
+                }
+                .boxed()
+            })
+            .await
+            .map_err(|err| err.to_safe_string())
     }
 }
 

golem-worker-executor/src/storage/keyvalue/mod.rs

@@ -14,6 +14,8 @@
 
 pub mod memory;
 pub mod multi_sqlite;
+pub mod namespace_routed;
+pub mod postgres;
 pub mod redis;
 pub mod sqlite;