refactor: refactor PgStore (#5309)

* refactor: refactor PgStore * fix: election use bytea and txn use Serializable to avoid read unrepeatable (#4) * fix: election use bytea as well * fix: use Serializable to avoid read unrepeatable * chore: remove unused error * ci: enable pg kvbackend and sqlness * ci: switch on pg_kvbackend feature * fix: fix sqlness runner * chore: add pg_kvbackend feature gate * build(ci): add feature gate * fix: add retry for `PgStore` txn * fix: correct `SET_IDLE_SESSION_TIMEOUT` --------- Co-authored-by: Yohan Wal <[email protected]> Co-authored-by: CookiePieWw <[email protected]>
GreptimeTeam · Jan 7, 2025 · 3d9df82 · 3d9df82
1 parent bc2f05d
commit 3d9df82
Show file tree

Hide file tree

Showing 18 changed files with 799 additions and 657 deletions.
diff --git a/.github/workflows/develop.yml b/.github/workflows/develop.yml
@@ -100,7 +100,7 @@ jobs:
       - name: Build greptime binaries
         shell: bash
         # `cargo gc` will invoke `cargo build` with specified args
-        run: cargo gc -- --bin greptime --bin sqlness-runner
+        run: cargo gc -- --bin greptime --bin sqlness-runner --features pg_kvbackend
       - name: Pack greptime binaries
         shell: bash
         run: |
@@ -261,7 +261,7 @@ jobs:
       - name: Build greptime bianry
         shell: bash
         # `cargo gc` will invoke `cargo build` with specified args
-        run: cargo gc --profile ci -- --bin greptime
+        run: cargo gc --profile ci -- --bin greptime --features pg_kvbackend
       - name: Pack greptime binary
         shell: bash
         run: |
@@ -573,6 +573,9 @@ jobs:
           - name: "Remote WAL"
             opts: "-w kafka -k 127.0.0.1:9092"
             kafka: true
+          - name: "Pg Kvbackend"
+            opts: "--setup-pg"
+            kafka: false
     timeout-minutes: 60
     steps:
       - uses: actions/checkout@v4

diff --git a/Cargo.lock b/Cargo.lock
diff --git a/Cargo.toml b/Cargo.toml
@@ -99,6 +99,7 @@ arrow-schema = { version = "51.0", features = ["serde"] }
 async-stream = "0.3"
 async-trait = "0.1"
 axum = { version = "0.6", features = ["headers"] }
+backon = "1"
 base64 = "0.21"
 bigdecimal = "0.4.2"
 bitflags = "2.4.1"

diff --git a/src/cli/src/bench.rs b/src/cli/src/bench.rs
@@ -78,7 +78,9 @@ impl BenchTableMetadataCommand {
         #[cfg(feature = "pg_kvbackend")]
         let kv_backend = if let Some(postgres_addr) = &self.postgres_addr {
             info!("Using postgres as kv backend");
-            PgStore::with_url(postgres_addr, 128).await.unwrap()
+            PgStore::with_url(postgres_addr, "greptime_metakv", 128)
+                .await
+                .unwrap()
         } else {
             kv_backend
         };

diff --git a/src/common/meta/Cargo.toml b/src/common/meta/Cargo.toml
@@ -6,7 +6,7 @@ license.workspace = true
 
 [features]
 testing = []
-pg_kvbackend = ["dep:tokio-postgres"]
+pg_kvbackend = ["dep:tokio-postgres", "dep:backon"]
 
 [lints]
 workspace = true
@@ -17,6 +17,7 @@ api.workspace = true
 async-recursion = "1.0"
 async-stream = "0.3"
 async-trait.workspace = true
+backon = { workspace = true, optional = true }
 base64.workspace = true
 bytes.workspace = true
 chrono.workspace = true

diff --git a/src/common/meta/src/error.rs b/src/common/meta/src/error.rs
@@ -639,15 +639,6 @@ pub enum Error {
         location: Location,
     },
 
-    #[snafu(display("Failed to parse {} from str to utf8", name))]
-    StrFromUtf8 {
-        name: String,
-        #[snafu(source)]
-        error: std::str::Utf8Error,
-        #[snafu(implicit)]
-        location: Location,
-    },
-
     #[snafu(display("Value not exists"))]
     ValueNotExist {
         #[snafu(implicit)]
@@ -658,8 +649,9 @@ pub enum Error {
     GetCache { source: Arc<Error> },
 
     #[cfg(feature = "pg_kvbackend")]
-    #[snafu(display("Failed to execute via Postgres"))]
+    #[snafu(display("Failed to execute via Postgres, sql: {}", sql))]
     PostgresExecution {
+        sql: String,
         #[snafu(source)]
         error: tokio_postgres::Error,
         #[snafu(implicit)]
@@ -693,6 +685,13 @@ pub enum Error {
         operation: String,
     },
 
+    #[cfg(feature = "pg_kvbackend")]
+    #[snafu(display("Postgres transaction retry failed"))]
+    PostgresTransactionRetryFailed {
+        #[snafu(implicit)]
+        location: Location,
+    },
+
     #[snafu(display(
         "Datanode table info not found, table id: {}, datanode id: {}",
         table_id,
@@ -756,8 +755,7 @@ impl ErrorExt for Error {
             | UnexpectedLogicalRouteTable { .. }
             | ProcedureOutput { .. }
             | FromUtf8 { .. }
-            | MetadataCorruption { .. }
-            | StrFromUtf8 { .. } => StatusCode::Unexpected,
+            | MetadataCorruption { .. } => StatusCode::Unexpected,
 
             SendMessage { .. } | GetKvCache { .. } | CacheNotGet { .. } => StatusCode::Internal,
 
@@ -807,7 +805,8 @@ impl ErrorExt for Error {
             PostgresExecution { .. }
             | CreatePostgresPool { .. }
             | GetPostgresConnection { .. }
-            | PostgresTransaction { .. } => StatusCode::Internal,
+            | PostgresTransaction { .. }
+            | PostgresTransactionRetryFailed { .. } => StatusCode::Internal,
             Error::DatanodeTableInfoNotFound { .. } => StatusCode::Internal,
         }
     }
@@ -818,6 +817,20 @@ impl ErrorExt for Error {
 }
 
 impl Error {
+    #[cfg(feature = "pg_kvbackend")]
+    /// Check if the error is a serialization error.
+    pub fn is_serialization_error(&self) -> bool {
+        match self {
+            Error::PostgresTransaction { error, .. } => {
+                error.code() == Some(&tokio_postgres::error::SqlState::T_R_SERIALIZATION_FAILURE)
+            }
+            Error::PostgresExecution { error, .. } => {
+                error.code() == Some(&tokio_postgres::error::SqlState::T_R_SERIALIZATION_FAILURE)
+            }
+            _ => false,
+        }
+    }
+
     /// Creates a new [Error::RetryLater] error from source `err`.
     pub fn retry_later<E: ErrorExt + Send + Sync + 'static>(err: E) -> Error {
         Error::RetryLater {

diff --git a/src/common/meta/src/kv_backend/etcd.rs b/src/common/meta/src/kv_backend/etcd.rs
@@ -591,7 +591,7 @@ mod tests {
     #[tokio::test]
     async fn test_range_2() {
         if let Some(kv_backend) = build_kv_backend().await {
-            test_kv_range_2_with_prefix(kv_backend, b"range2/".to_vec()).await;
+            test_kv_range_2_with_prefix(&kv_backend, b"range2/".to_vec()).await;
         }
     }
 
@@ -618,7 +618,8 @@ mod tests {
         if let Some(kv_backend) = build_kv_backend().await {
             let prefix = b"deleteRange/";
             prepare_kv_with_prefix(&kv_backend, prefix.to_vec()).await;
-            test_kv_delete_range_with_prefix(kv_backend, prefix.to_vec()).await;
+            test_kv_delete_range_with_prefix(&kv_backend, prefix.to_vec()).await;
+            unprepare_kv(&kv_backend, prefix).await;
         }
     }
 
@@ -627,20 +628,20 @@ mod tests {
         if let Some(kv_backend) = build_kv_backend().await {
             let prefix = b"batchDelete/";
             prepare_kv_with_prefix(&kv_backend, prefix.to_vec()).await;
-            test_kv_batch_delete_with_prefix(kv_backend, prefix.to_vec()).await;
+            test_kv_batch_delete_with_prefix(&kv_backend, prefix.to_vec()).await;
+            unprepare_kv(&kv_backend, prefix).await;
         }
     }
 
     #[tokio::test]
     async fn test_etcd_txn() {
         if let Some(kv_backend) = build_kv_backend().await {
-            let kv_backend_ref = Arc::new(kv_backend);
-            test_txn_one_compare_op(kv_backend_ref.clone()).await;
-            text_txn_multi_compare_op(kv_backend_ref.clone()).await;
-            test_txn_compare_equal(kv_backend_ref.clone()).await;
-            test_txn_compare_greater(kv_backend_ref.clone()).await;
-            test_txn_compare_less(kv_backend_ref.clone()).await;
-            test_txn_compare_not_equal(kv_backend_ref).await;
+            test_txn_one_compare_op(&kv_backend).await;
+            text_txn_multi_compare_op(&kv_backend).await;
+            test_txn_compare_equal(&kv_backend).await;
+            test_txn_compare_greater(&kv_backend).await;
+            test_txn_compare_less(&kv_backend).await;
+            test_txn_compare_not_equal(&kv_backend).await;
         }
     }
 }
diff --git a/src/common/meta/src/kv_backend/memory.rs b/src/common/meta/src/kv_backend/memory.rs
@@ -355,7 +355,7 @@ mod tests {
     async fn test_range_2() {
         let kv = MemoryKvBackend::<Error>::new();
 
-        test_kv_range_2(kv).await;
+        test_kv_range_2(&kv).await;
     }
 
     #[tokio::test]
@@ -376,24 +376,24 @@ mod tests {
     async fn test_delete_range() {
         let kv_backend = mock_mem_store_with_data().await;
 
-        test_kv_delete_range(kv_backend).await;
+        test_kv_delete_range(&kv_backend).await;
     }
 
     #[tokio::test]
     async fn test_batch_delete() {
         let kv_backend = mock_mem_store_with_data().await;
 
-        test_kv_batch_delete(kv_backend).await;
+        test_kv_batch_delete(&kv_backend).await;
     }
 
     #[tokio::test]
     async fn test_memory_txn() {
-        let kv_backend = Arc::new(MemoryKvBackend::<Error>::new());
-        test_txn_one_compare_op(kv_backend.clone()).await;
-        text_txn_multi_compare_op(kv_backend.clone()).await;
-        test_txn_compare_equal(kv_backend.clone()).await;
-        test_txn_compare_greater(kv_backend.clone()).await;
-        test_txn_compare_less(kv_backend.clone()).await;
-        test_txn_compare_not_equal(kv_backend).await;
+        let kv_backend = MemoryKvBackend::<Error>::new();
+        test_txn_one_compare_op(&kv_backend).await;
+        text_txn_multi_compare_op(&kv_backend).await;
+        test_txn_compare_equal(&kv_backend).await;
+        test_txn_compare_greater(&kv_backend).await;
+        test_txn_compare_less(&kv_backend).await;
+        test_txn_compare_not_equal(&kv_backend).await;
     }
 }