Use shared hashmap

Author: kovmir

pg_sqlwall.c

@@ -1,11 +1,11 @@
 #include "postgres.h"
 #include "fmgr.h"
-#include "nodes/nodes.h"
+
+#include "funcapi.h"
 #include "miscadmin.h"
 #include "storage/ipc.h"
-#include "storage/shmem.h"
 #include "storage/lwlock.h"
-#include "utils/builtins.h"
+#include "storage/shmem.h"
 
 /* Constants and Macros */
 PG_MODULE_MAGIC;
@@ -16,26 +16,31 @@ PG_MODULE_MAGIC;
 #endif
 
 /* Is there a better size candidate? */
-#define DENY_LIST_SIZE 128
-/* CMD_NOTHING tells us about enum size as it is the last enum member. */
+#define DENY_LIST_MAX_SIZE 128
+/* CMD_NOTHING tells us about the enum size as it is the last enum member. */
 #define CMD_TYPE_COUNT CMD_NOTHING+1
 
-typedef struct pgswUserPerm {
+typedef struct pgswHashKey /* Hash entry key. */
+{
 	Oid userId;
-	/* Lists SQL command types the user is NOT allowed to run.
+} pgswHashKey;
+
+typedef struct pgswHashEntry
+{
+	slock_t mutex; /* Protects the entry. */
+	Oid userId;    /* User OID. */
+
+	/* Lists SQL command types the user is not allowed to run.
 	 * CmdType enum members are indexes for this array, i.e.
 	 * denyCmd[CMD_DELETE] = true to prohibit SQL DELETE statements. */
 	bool denyCmd[CMD_TYPE_COUNT];
-} pgswUserPerm;
-
-typedef struct pgswSharedState {
-	LWLock* lock;
-	pgswUserPerm denyList[DENY_LIST_SIZE];
-	int denyListSize;
-} pgswSharedState;
+} pgswHashEntry;
 
 /* Function Prototypes */
 void _PG_init(void);
+static uint32 gen_hash_key(const void *key, Size keysize);
+static int compare_hash_key(const void *key1, const void *key2, Size keysize);
+
 static void pgsw_shmem_request(void);
 static void pgsw_shmem_startup(void);
 
@@ -45,7 +50,7 @@ PG_FUNCTION_INFO_V1(my_function);
 /* Global variables */
 static shmem_request_hook_type prev_shmem_request_hook = NULL;
 static shmem_startup_hook_type prev_shmem_startup_hook = NULL;
-pgswSharedState *pgsw_shmem = NULL;
+static HTAB *pgsw_hash = NULL;
 
 void
 _PG_init(void)
@@ -66,26 +71,45 @@ pgsw_shmem_request(void)
 	if(prev_shmem_request_hook)
 		prev_shmem_request_hook();
 
-	RequestAddinShmemSpace(MAXALIGN(sizeof(pgswSharedState)));
+	RequestAddinShmemSpace(
+			hash_estimate_size(DENY_LIST_MAX_SIZE, sizeof(pgswHashEntry)) );
 	RequestNamedLWLockTranche("pg_sqlwall", 1);
 }
 
+int
+compare_hash_key(const void *key1, const void *key2, Size keysize)
+{
+	const pgswHashKey *k1 = (const pgswHashKey *) key1;
+	const pgswHashKey *k2 = (const pgswHashKey *) key2;
+	return (k1->userId == k2->userId) ? 0 : 1;
+}
+
+uint32
+gen_hash_key(const void *key, Size keysize)
+{
+	const pgswHashKey *k = (const pgswHashKey *) key;
+	return (uint32) k->userId;
+}
+
 void
 pgsw_shmem_startup(void)
 {
-	bool found;
+	HASHCTL info;
 
 	if(prev_shmem_startup_hook)
 		prev_shmem_startup_hook();
 
 	LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE);
-
-	pgsw_shmem = ShmemInitStruct("pg_sqlwall", sizeof(pgswSharedState), &found);
-	if(!found) {
-		pgsw_shmem->denyListSize = 0;
-		pgsw_shmem->lock = &(GetNamedLWLockTranche("pg_sqlwall"))->lock;
-	}
-
+	memset(&info, 0, sizeof(info));
+	info.keysize = sizeof(pgswHashKey);
+	info.entrysize = sizeof(pgswHashKey);
+	info.hash = gen_hash_key;
+	info.match = compare_hash_key;
+	pgsw_hash = ShmemInitHash("pg_sqlwall hash",
+								MaxConnections,
+								MaxConnections,
+								&info,
+								HASH_ELEM|HASH_FUNCTION|HASH_COMPARE);
 	LWLockRelease(AddinShmemInitLock);
 }