Shared memory hash table

pgrx-examples/shmem/Cargo.toml

@@ -17,7 +17,7 @@ edition = "2021"
 crate-type = ["cdylib"]
 
 [features]
-default = ["pg13"]
+default = ["pg13", "pgrx/cshim"]
 pg12 = ["pgrx/pg12", "pgrx-tests/pg12" ]
 pg13 = ["pgrx/pg13", "pgrx-tests/pg13" ]
 pg14 = ["pgrx/pg14", "pgrx-tests/pg14" ]

pgrx-examples/shmem/src/lib.rs

@@ -11,6 +11,7 @@ use pgrx::atomics::*;
 use pgrx::lwlock::PgLwLock;
 use pgrx::prelude::*;
 use pgrx::shmem::*;
+use pgrx::shmem_hash::*;
 use pgrx::{pg_shmem_init, warning};
 use serde::*;
 use std::iter::Iterator;
@@ -39,6 +40,7 @@ static HASH: PgLwLock<heapless::FnvIndexMap<i32, i32, 4>> = PgLwLock::new();
 static STRUCT: PgLwLock<Pgtest> = PgLwLock::new();
 static PRIMITIVE: PgLwLock<i32> = PgLwLock::new();
 static ATOMIC: PgAtomic<std::sync::atomic::AtomicBool> = PgAtomic::new();
+static HASH_TABLE: PgHashMap<i64, i64> = PgHashMap::new(250);
 
 #[pg_guard]
 pub extern "C" fn _PG_init() {
@@ -48,6 +50,7 @@ pub extern "C" fn _PG_init() {
     pg_shmem_init!(STRUCT);
     pg_shmem_init!(PRIMITIVE);
     pg_shmem_init!(ATOMIC);
+    pg_shmem_init!(HASH_TABLE);
 }
 
 #[pg_extern]
@@ -60,6 +63,26 @@ fn vec_count() -> i32 {
     VEC.share().len() as i32
 }
 
+#[pg_extern]
+fn hash_table_insert(key: i64, value: i64) -> Option<i64> {
+    HASH_TABLE.insert(key, value).unwrap()
+}
+
+#[pg_extern]
+fn hash_table_get(key: i64) -> Option<i64> {
+    HASH_TABLE.get(key)
+}
+
+#[pg_extern]
+fn hash_table_remove(key: i64) -> Option<i64> {
+    HASH_TABLE.remove(key)
+}
+
+#[pg_extern]
+fn hash_table_len() -> i64 {
+    HASH_TABLE.len() as i64
+}
+
 #[pg_extern]
 fn vec_drain() -> SetOfIterator<'static, Pgtest> {
     let mut vec = VEC.exclusive();

pgrx-tests/src/tests/shmem_tests.rs

@@ -8,25 +8,28 @@
 //LICENSE
 //LICENSE Use of this source code is governed by the MIT license that can be found in the LICENSE file.
 use pgrx::prelude::*;
-use pgrx::{pg_shmem_init, PgAtomic, PgLwLock, PgSharedMemoryInitialization};
+use pgrx::{pg_shmem_init, PgAtomic, PgHashMap, PgLwLock, PgSharedMemoryInitialization};
 use std::sync::atomic::AtomicBool;
 
 static ATOMIC: PgAtomic<AtomicBool> = PgAtomic::new();
 static LWLOCK: PgLwLock<bool> = PgLwLock::new();
+static HASH_MAP: PgHashMap<i64, i64> = PgHashMap::new(500);
 
 #[pg_guard]
 pub extern "C" fn _PG_init() {
     // This ensures that this functionality works across PostgreSQL versions
     pg_shmem_init!(ATOMIC);
     pg_shmem_init!(LWLOCK);
+    pg_shmem_init!(HASH_MAP);
 }
+
 #[cfg(any(test, feature = "pg_test"))]
 #[pgrx::pg_schema]
 mod tests {
     #[allow(unused_imports)]
     use crate as pgrx_tests;
 
-    use crate::tests::shmem_tests::LWLOCK;
+    use crate::tests::shmem_tests::{HASH_MAP, LWLOCK};
     use pgrx::prelude::*;
 
     #[pg_test]
@@ -53,4 +56,54 @@ mod tests {
         });
         let _lock = LWLOCK.exclusive();
     }
+
+    #[pg_test]
+    pub fn test_pg_hash_map() {
+        use rand::prelude::IteratorRandom;
+
+        for i in 1..250 {
+            assert_eq!(HASH_MAP.insert(i, i), Ok(None));
+        }
+
+        assert_eq!(HASH_MAP.len(), 249);
+
+        for i in 1..250 {
+            assert_eq!(HASH_MAP.get(i), Some(i));
+        }
+
+        assert_eq!(HASH_MAP.len(), 249);
+
+        for i in 251..500 {
+            assert_eq!(HASH_MAP.get(i), None);
+        }
+
+        assert_eq!(HASH_MAP.len(), 249);
+
+        for i in 1..250 {
+            assert_eq!(HASH_MAP.insert(i, i), Ok(Some(i)));
+        }
+
+        assert_eq!(HASH_MAP.len(), 249);
+
+        for i in 1..250 {
+            assert_eq!(HASH_MAP.remove(i), Some(i));
+        }
+
+        assert_eq!(HASH_MAP.len(), 0);
+
+        for i in 1..250 {
+            assert_eq!(HASH_MAP.get(i), None);
+        }
+
+        assert_eq!(HASH_MAP.len(), 0);
+
+        for _ in 0..25_000 {
+            for key in 0..250 {
+                let value = (0..1000).choose(&mut rand::thread_rng()).unwrap();
+                assert!(HASH_MAP.insert(key, value).is_ok());
+            }
+        }
+
+        assert_eq!(HASH_MAP.len(), 250);
+    }
 }

pgrx/src/lib.rs

@@ -65,6 +65,8 @@ pub mod pg_catalog;
 pub mod pgbox;
 pub mod rel;
 pub mod shmem;
+#[cfg(feature = "cshim")]
+pub mod shmem_hash;
 pub mod spi;
 #[cfg(feature = "cshim")]
 pub mod spinlock;
@@ -106,6 +108,8 @@ pub use nodes::*;
 pub use pgbox::*;
 pub use rel::*;
 pub use shmem::*;
+#[cfg(feature = "cshim")]
+pub use shmem_hash::*;
 pub use spi::Spi; // only Spi.  We don't want the top-level namespace polluted with spi::Result and spi::Error
 pub use stringinfo::*;
 pub use trigger_support::*;