WIP diesel for persistence

Cargo.lock

@@ -264,6 +264,7 @@ dependencies = [
  "camino",
  "clap",
  "control_plane",
+ "diesel",
  "futures",
  "hyper",
  "pageserver_api",
@@ -1605,6 +1606,39 @@ dependencies = [
  "rusticata-macros",
 ]
 
+[[package]]
+name = "diesel"
+version = "2.1.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "62c6fcf842f17f8c78ecf7c81d75c5ce84436b41ee07e03f490fbb5f5a8731d8"
+dependencies = [
+ "diesel_derives",
+ "libsqlite3-sys",
+ "serde_json",
+ "time",
+]
+
+[[package]]
+name = "diesel_derives"
+version = "2.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "ef8337737574f55a468005a83499da720f20c65586241ffea339db9ecdfd2b44"
+dependencies = [
+ "diesel_table_macro_syntax",
+ "proc-macro2",
+ "quote",
+ "syn 2.0.32",
+]
+
+[[package]]
+name = "diesel_table_macro_syntax"
+version = "0.1.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "fc5557efc453706fed5e4fa85006fe9817c224c3f480a34c7e5959fd700921c5"
+dependencies = [
+ "syn 2.0.32",
+]
+
 [[package]]
 name = "digest"
 version = "0.10.7"
@@ -2623,6 +2657,16 @@ dependencies = [
  "winapi",
 ]
 
+[[package]]
+name = "libsqlite3-sys"
+version = "0.27.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cf4e226dcd58b4be396f7bd3c20da8fdee2911400705297ba7d2d7cc2c30f716"
+dependencies = [
+ "pkg-config",
+ "vcpkg",
+]
+
 [[package]]
 name = "linux-raw-sys"
 version = "0.1.4"

control_plane/attachment_service/Cargo.toml

@@ -20,6 +20,8 @@ tokio.workspace = true
 tokio-util.workspace = true
 tracing.workspace = true
 
+diesel = { version = "2.1.4", features = ["sqlite", "serde_json"] }
+
 utils = { path = "../../libs/utils/" }
 control_plane = { path = ".." }
 workspace_hack = { version = "0.1", path = "../../workspace_hack" }

control_plane/attachment_service/migrations/2024-01-07-211257_create_tenant_shards/down.sql

@@ -0,0 +1 @@
+DROP TABLE tenant_shards;

control_plane/attachment_service/migrations/2024-01-07-211257_create_tenant_shards/up.sql

@@ -0,0 +1,11 @@
+CREATE TABLE tenant_shards (
+  id INTEGER PRIMARY KEY NOT NULL,
+  tenant_id VARCHAR NOT NULL,
+  shard_number INTEGER NOT NULL,
+  shard_count INTEGER NOT NULL,
+  shard_stripe_size INTEGER NOT NULL,
+  generation INTEGER NOT NULL,
+  placement_policy VARCHAR NOT NULL,
+  -- config is JSON encoded, opaque to the database.
+  config TEXT NOT NULL
+);

control_plane/attachment_service/migrations/2024-01-07-212945_create_nodes/down.sql

@@ -0,0 +1 @@
+DROP TABLE nodes;

control_plane/attachment_service/migrations/2024-01-07-212945_create_nodes/up.sql

@@ -0,0 +1,10 @@
+CREATE TABLE nodes (
+  node_id BIGINT PRIMARY KEY NOT NULL,
+
+  scheduling_policy VARCHAR NOT NULL,
+
+  listen_http_addr VARCHAR NOT NULL,
+  listen_http_port INTEGER NOT NULL,
+  listen_pg_addr VARCHAR NOT NULL,
+  listen_pg_port INTEGER NOT NULL
+);

control_plane/attachment_service/src/lib.rs

@@ -3,8 +3,10 @@ use utils::seqwait::MonotonicCounter;
 mod compute_hook;
 pub mod http;
 mod node;
+mod persistence;
 mod reconciler;
 mod scheduler;
+mod schema;
 pub mod service;
 mod tenant_state;
 

control_plane/attachment_service/src/node.rs

@@ -1,4 +1,6 @@
-use control_plane::attachment_service::{NodeAvailability, NodeSchedulingPolicy};
+use control_plane::attachment_service::NodeAvailability;
+use diesel::expression::AsExpression;
+use serde::{Deserialize, Serialize};
 use utils::id::NodeId;
 
 #[derive(Clone)]
@@ -15,6 +17,22 @@ pub(crate) struct Node {
     pub(crate) listen_pg_port: u16,
 }
 
+#[derive(Serialize, Deserialize, Clone, Copy, Debug, AsExpression)]
+#[diesel(sql_type = diesel::sql_types::VarChar)]
+pub enum NodeSchedulingPolicy {
+    // Normal, happy state
+    Active,
+
+    // A newly added node: gradually move some work here.
+    Filling,
+
+    // Do not schedule new work here, but leave configured locations in place.
+    Pause,
+
+    // Do not schedule work here.  Gracefully move work away, as resources allow.
+    Draining,
+}
+
 impl Node {
     pub(crate) fn base_url(&self) -> String {
         format!("http://{}:{}", self.listen_http_addr, self.listen_http_port)

control_plane/attachment_service/src/persistence.rs

@@ -1,13 +1,72 @@
+use std::env;
+
+use anyhow::Context;
+use diesel::prelude::*;
+use diesel::{Connection, SqliteConnection};
+use utils::generation::Generation;
+
+use crate::node::NodeSchedulingPolicy;
+use crate::PlacementPolicy;
+
 /// The attachment service does not store most of its state durably.
 ///
 /// The essential things to store durably are:
 /// - generation numbers, as these must always advance monotonically to ensure data safety.
-/// - PlacementPolicy and TenantConfig, as these are set externally.
-struct Persistence {}
-
-/// Parts of TenantState that are stored durably
-struct TenantPersistence {
-    pub(crate) generation: Generation,
-    pub(crate) policy: PlacementPolicy,
-    pub(crate) config: TenantConfig,
+/// - Tenant's PlacementPolicy and TenantConfig, as the source of truth for these is something external.
+/// - Node's scheduling policies, as the source of truth for these is something external.
+///
+/// Other things we store durably as an implementation detail:
+/// - Node's host/port: this could be avoided it we made nodes emit a self-registering heartbeat,
+///   but it is operationally simpler to make this service the authority for which nodes
+///   it talks to.
+struct Persistence {
+    database_url: String,
+}
+
+impl Persistence {
+    fn new() -> Self {
+        let database_url = env::var("DATABASE_URL").expect("DATABASE_URL must be set");
+        Self { database_url }
+    }
+
+    async fn insert_node(&self, node: NodePersistence) -> anyhow::Result<()> {
+        let conn = SqliteConnection::establish(&self.database_url).context("Opening database")?;
+        diesel::insert_into(crate::schema::nodes::table)
+            .values(&node)
+            .get_result(&mut conn)
+            .into()
+    }
+
+    async fn list_nodes(&self) -> anyhow::Result<Vec<NodePersistence>> {
+        let mut conn =
+            SqliteConnection::establish(&self.database_url).context("Opening database")?;
+
+        crate::schema::nodes::dsl::nodes
+            .select(NodePersistence::as_select())
+            .load(&mut conn)
+            .into()
+    }
+}
+
+/// Parts of [`crate::tenant_state::TenantState`] that are stored durably
+#[derive(Selectable)]
+#[diesel(table_name = crate::schema::tenant_shards)]
+pub(crate) struct TenantShardPersistence {
+    pub(crate) generation: Generation,
+    pub(crate) placement_policy: PlacementPolicy,
+    pub(crate) config: serde_json::Value,
+}
+
+/// Parts of [`crate::node::Node`] that are stored durably
+#[derive(Selectable, Insertable)]
+#[diesel(table_name = crate::schema::nodes)]
+pub(crate) struct NodePersistence {
+    pub(crate) node_id: i64,
+
+    pub(crate) scheduling_policy: NodeSchedulingPolicy,
+    pub(crate) listen_http_addr: String,
+    pub(crate) listen_http_port: i32,
+
+    pub(crate) listen_pg_addr: String,
+    pub(crate) listen_pg_port: i32,
 }

control_plane/attachment_service/src/schema.rs

@@ -0,0 +1,27 @@
+// @generated automatically by Diesel CLI.
+
+diesel::table! {
+    nodes (node_id) {
+        node_id -> BigInt,
+        scheduling_policy -> Text,
+        listen_http_addr -> Text,
+        listen_http_port -> Integer,
+        listen_pg_addr -> Text,
+        listen_pg_port -> Integer,
+    }
+}
+
+diesel::table! {
+    tenant_shards (id) {
+        id -> Integer,
+        tenant_id -> Text,
+        shard_number -> Integer,
+        shard_count -> Integer,
+        shard_stripe_size -> Integer,
+        generation -> Integer,
+        placement_policy -> Text,
+        config -> Text,
+    }
+}
+
+diesel::allow_tables_to_appear_in_same_query!(nodes, tenant_shards,);

control_plane/attachment_service/src/tenant_state.rs

@@ -26,6 +26,9 @@ pub(crate) struct TenantState {
 
     pub(crate) shard: ShardIdentity,
 
+    // Runtime only: sequence used to coordinate when updating this object while
+    // with background reconcilers may be running.  A reconciler runs to a particular
+    // sequence.
     pub(crate) sequence: Sequence,
 
     // Latest generation number: next time we attach, increment this
@@ -45,6 +48,8 @@ pub(crate) struct TenantState {
     // with `Self::reconcile`.
     pub(crate) observed: ObservedState,
 
+    // Tenant configuration, passed through opaquely to the pageserver.  Identical
+    // for all shards in a tenant.
     pub(crate) config: TenantConfig,
 
     /// If a reconcile task is currently in flight, it may be joined here (it is

control_plane/src/attachment_service.rs

@@ -127,18 +127,13 @@ impl FromStr for NodeAvailability {
     }
 }
 
+/// FIXME: this is a duplicate of the type in the attachment_service crate, because the
+/// type needs to be defined with diesel traits in there.
 #[derive(Serialize, Deserialize, Clone, Copy)]
 pub enum NodeSchedulingPolicy {
-    // Normal, happy state
     Active,
-
-    // A newly added node: gradually move some work here.
     Filling,
-
-    // Do not schedule new work here, but leave configured locations in place.
     Pause,
-
-    // Do not schedule work here.  Gracefully move work away, as resources allow.
     Draining,
 }
 

diesel.toml

@@ -0,0 +1,9 @@
+# For documentation on how to configure this file,
+# see https://diesel.rs/guides/configuring-diesel-cli
+
+[print_schema]
+file = "control_plane/attachment_service/src/schema.rs"
+custom_type_derives = ["diesel::query_builder::QueryId"]
+
+[migrations_directory]
+dir = "control_plane/attachment_service/migrations"