diff --git a/src/test/modules/simple_cache/Makefile b/src/test/modules/simple_cache/Makefile new file mode 100644 index 0000000..b9a0b28 --- /dev/null +++ b/src/test/modules/simple_cache/Makefile @@ -0,0 +1,18 @@ +# src/test/modules/simple_cache/Makefile + +MODULES = simple_cache + +EXTENSION = simple_cache +DATA = simple_cache--1.0.sql +PGFILEDESC = "simple_cache -- a simple cache extension" + +ifdef USE_PGXS +PG_CONFIG = pg_config +PGXS := $(shell $(PG_CONFIG) --pgxs) +include $(PGXS) +else +subdir = src/test/modules/simple_cache +top_builddir = ../../../.. +include $(top_builddir)/src/Makefile.global +include $(top_srcdir)/contrib/contrib-global.mk +endif diff --git a/src/test/modules/simple_cache/simple_cache--1.0.sql b/src/test/modules/simple_cache/simple_cache--1.0.sql new file mode 100644 index 0000000..d5ce5fa --- /dev/null +++ b/src/test/modules/simple_cache/simple_cache--1.0.sql @@ -0,0 +1,39 @@ +/* src/test/modules/simple_cache/simple_cache--1.0.sql */ + +-- complain if script is sourced in psql, rather than via CREATE EXTENSION +\echo Use "CREATE EXTENSION simple_cache" to load this file. \quit + +CREATE FUNCTION simple_cache_put(key text, value text) +RETURNS boolean +AS 'MODULE_PATHNAME' +LANGUAGE C; + +CREATE FUNCTION simple_cache_get(key text) +RETURNS pg_catalog.text +AS 'MODULE_PATHNAME' +LANGUAGE C; + +CREATE FUNCTION simple_cache_drop(key text) +RETURNS boolean +AS 'MODULE_PATHNAME' +LANGUAGE C; + +CREATE FUNCTION simple_cache_list(out key text, out value text) +RETURNS setof record +AS 'MODULE_PATHNAME' +LANGUAGE C; + +CREATE FUNCTION simple_cache_set_size_limit(n integer) +RETURNS VOID +AS 'MODULE_PATHNAME' +LANGUAGE C; + +CREATE FUNCTION simple_cache_dump_area() +RETURNS VOID +AS 'MODULE_PATHNAME' +LANGUAGE C; + +CREATE FUNCTION simple_cache_dump_hash_table() +RETURNS VOID +AS 'MODULE_PATHNAME' +LANGUAGE C; diff --git a/src/test/modules/simple_cache/simple_cache.c b/src/test/modules/simple_cache/simple_cache.c new file mode 100644 index 0000000..7ea149d --- /dev/null +++ b/src/test/modules/simple_cache/simple_cache.c @@ -0,0 +1,393 @@ +/* ------------------------------------------------------------------------- + * + * simple_cache.c + * A toy module to demonstrate DHT hash tables. + * + * Copyright (c) 2016, PostgreSQL Global Development Group + * + * IDENTIFICATION + * src/test/modules/simple_cache/simple_cache.c + * + * ------------------------------------------------------------------------- + */ + +#include "postgres.h" + +#include "fmgr.h" +#include "funcapi.h" +#include "miscadmin.h" +#include "nodes/execnodes.h" +#include "storage/dht.h" +#include "storage/dsa.h" +#include "storage/ipc.h" +#include "storage/lwlock.h" +#include "storage/shmem.h" +#include "utils/builtins.h" +#include "utils/memutils.h" +#include "utils/hsearch.h" /* for tag_hash */ + +PG_MODULE_MAGIC; + +PG_FUNCTION_INFO_V1(simple_cache_put); +PG_FUNCTION_INFO_V1(simple_cache_get); +PG_FUNCTION_INFO_V1(simple_cache_drop); +PG_FUNCTION_INFO_V1(simple_cache_list); +PG_FUNCTION_INFO_V1(simple_cache_set_size_limit); +PG_FUNCTION_INFO_V1(simple_cache_dump_area); +PG_FUNCTION_INFO_V1(simple_cache_dump_hash_table); + +void _PG_init(void); + +/* + * The shmem struct we use to give dynamic shared area and hash table handles + * to other processes. + */ +typedef struct +{ + /* The handle for attaching to the shared area. */ + dsa_handle area_handle; + /* The handle for attaching to the hash table. */ + dht_hash_table_handle hash_table_handle; +} simple_cache_shared; + +/* + * The state we need for each backend. + */ +typedef struct +{ + /* The shared memory area. */ + dsa_area *area; + /* The hash table. */ + dht_hash_table *hash_table; +} simple_cache_state; + +#define SIMPLE_CACHE_KEY_SIZE 48 + +typedef struct +{ + char key[SIMPLE_CACHE_KEY_SIZE]; + Size value_size; + dsa_pointer value; +} simple_cache_entry; + +void +_PG_init(void) +{ +} + +static simple_cache_state * +get_simple_cache_state(void) +{ + static simple_cache_state result = { 0 }; + + if (result.hash_table == NULL) + { + dht_parameters params = { + SIMPLE_CACHE_KEY_SIZE, + sizeof(simple_cache_entry), + memcmp, + tag_hash, + 0, /* assigned below */ + "dht lock tranche" + }; + simple_cache_shared *shared; + bool found; + + /* Create or attach to the memory area and hash table on demand. */ + LWLockAcquire(AddinShmemInitLock, LW_EXCLUSIVE); + shared = ShmemInitStruct("simple_cache", sizeof(simple_cache_shared), &found); + if (!found) + { + /* First caller. Set up 'not yet created' state. */ + shared->area_handle = 0; + shared->hash_table_handle = 0; + } + /* Create shared area if not yet created. */ + if (shared->area_handle == 0) + { + MemoryContext old_context; + + /* + * Here, and in several places below, we switch to + * TopMemoryContext because the backend-local dsa_area + * and dht_hash_table objects are palloc'd and we want them + * to stick around until backend exit. All of this jiggery-pokery + * and pinning is required because we're creating objects that + * we want to use until cluster shutdown. None of that stuff + * is necessary when using DSA or DHT for parallel execution: + * in that context we want the automatic cleanup that we're + * effectively disabling here. + */ + old_context = MemoryContextSwitchTo(TopMemoryContext); + result.area = dsa_create_dynamic(LWLockNewTrancheId(), + "simple_cache dsa_area"); + MemoryContextSwitchTo(old_context); + if (result.area == NULL) + elog(ERROR, "simple_cache: could not create area"); + + /* + * We want this area and its memory to survive even when + * there are no backends attached. + */ + dsa_pin(result.area); + /* + * We want this area's memory to stay mapped in, regardless of + * what happens to the current resource owner. + */ + dsa_pin_mapping(result.area); + /* We want other backends to be able to attach. */ + shared->area_handle = dsa_get_handle(result.area); + } + /* Create hash table if not yet created. */ + if (shared->hash_table_handle == 0) + { + MemoryContext old_context; + + params.tranche_id = LWLockNewTrancheId(); + old_context = MemoryContextSwitchTo(TopMemoryContext); + result.hash_table = dht_create(result.area, ¶ms); + MemoryContextSwitchTo(old_context); + if (result.hash_table == NULL) + elog(ERROR, "simple_cache: could not create hash table"); + shared->hash_table_handle = dht_get_hash_table_handle(result.hash_table); + } + /* Attach to shared area if not yet attached. */ + if (result.area == NULL) + { + MemoryContext old_context; + + old_context = MemoryContextSwitchTo(TopMemoryContext); + result.area = dsa_attach_dynamic(shared->area_handle); + MemoryContextSwitchTo(old_context); + + dsa_pin_mapping(result.area); + } + /* Attach to hash table if not yet attached. */ + if (result.hash_table == NULL) + { + MemoryContext old_context; + + old_context = MemoryContextSwitchTo(TopMemoryContext); + result.hash_table = dht_attach(result.area, ¶ms, + shared->hash_table_handle); + MemoryContextSwitchTo(old_context); + } + LWLockRelease(AddinShmemInitLock); + } + + return &result; +} + +Datum +simple_cache_put(PG_FUNCTION_ARGS) +{ + bool found; + char padded_key[SIMPLE_CACHE_KEY_SIZE]; + text *key; + text *value; + simple_cache_state *state; + simple_cache_entry *entry; + dsa_pointer new_value; + + if (PG_ARGISNULL(0) || PG_ARGISNULL(1)) + elog(ERROR, "simple_cache_put: arguments must be non-NULL"); + + key = PG_GETARG_TEXT_PP(0); + value = PG_GETARG_TEXT_PP(1); + + /* Build a NUL-padded version of the key. */ + if (VARSIZE_ANY_EXHDR(key) + 1 > SIMPLE_CACHE_KEY_SIZE) + elog(ERROR, "simple_cache_put: key too long"); + memset(padded_key, 0, sizeof(padded_key)); + memcpy(padded_key, VARDATA_ANY(key), VARSIZE_ANY_EXHDR(key)); + + state = get_simple_cache_state(); + + /* Find or create an exclusively locked hash table entry. */ + entry = dht_find_or_insert(state->hash_table, padded_key, &found); + if (entry == NULL) + elog(ERROR, "simple_cache_put: out of memory"); + + /* Copy the value into DSA memory. */ + new_value = dsa_allocate(state->area, VARSIZE_ANY_EXHDR(value)); + if (!DsaPointerIsValid(new_value)) + { + /* + * No memory. Delete this entry if we just created it, or just unlock + * if it was already there, and bail out. + */ + if (!found) + dht_delete_entry(state->hash_table, entry); + else + dht_release(state->hash_table, entry); + elog(ERROR, "simple_cache_put: out of memory"); + } + memcpy(dsa_get_address(state->area, new_value), + VARDATA_ANY(value), + VARSIZE_ANY_EXHDR(value)); + /* Put it into the entry (freeing the old value first if needed). */ + if (found) + dsa_free(state->area, entry->value); + entry->value = new_value; + entry->value_size = VARSIZE_ANY_EXHDR(value); + + /* Unlock the hash table entry. */ + dht_release(state->hash_table, entry); + + PG_RETURN_BOOL(found); +} + +Datum +simple_cache_get(PG_FUNCTION_ARGS) +{ + char padded_key[SIMPLE_CACHE_KEY_SIZE]; + text *key; + text *value; + simple_cache_state *state; + simple_cache_entry *entry; + + /* Reject invalid keys with NULL. */ + if (PG_ARGISNULL(0)) + PG_RETURN_NULL(); + key = PG_GETARG_TEXT_PP(0); + if (VARSIZE_ANY_EXHDR(key) + 1 > SIMPLE_CACHE_KEY_SIZE) + PG_RETURN_NULL(); + + /* Build a NUL-padded version of the key. */ + memset(padded_key, 0, sizeof(padded_key)); + memcpy(padded_key, VARDATA_ANY(key), VARSIZE_ANY_EXHDR(key)); + + state = get_simple_cache_state(); + + /* Try to find the key. */ + entry = dht_find(state->hash_table, padded_key, false); + if (entry == NULL) + PG_RETURN_NULL(); + value = palloc(entry->value_size + VARHDRSZ); + SET_VARSIZE(value, entry->value_size + VARHDRSZ); + memcpy(VARDATA_ANY(value), + dsa_get_address(state->area, entry->value), + entry->value_size); + dht_release(state->hash_table, entry); + + PG_RETURN_TEXT_P(value); +} + +Datum +simple_cache_drop(PG_FUNCTION_ARGS) +{ + char padded_key[SIMPLE_CACHE_KEY_SIZE]; + text *key; + simple_cache_state *state; + simple_cache_entry *entry; + + /* Reject invalid keys with NULL. */ + if (PG_ARGISNULL(0)) + PG_RETURN_NULL(); + key = PG_GETARG_TEXT_PP(0); + if (VARSIZE_ANY_EXHDR(key) + 1 > SIMPLE_CACHE_KEY_SIZE) + PG_RETURN_NULL(); + + /* Build a NUL-padded version of the key. */ + memset(padded_key, 0, sizeof(padded_key)); + memcpy(padded_key, VARDATA_ANY(key), VARSIZE_ANY_EXHDR(key)); + + state = get_simple_cache_state(); + + /* Find it and lock it exclusively (prerequisite for deleting). */ + entry = dht_find(state->hash_table, padded_key, true); + if (entry != NULL) + { + /* Free the value that we own in DSA memory. */ + dsa_free(state->area, entry->value); + dht_delete_entry(state->hash_table, entry); + } + PG_RETURN_BOOL(entry != NULL); +} + +Datum +simple_cache_list(PG_FUNCTION_ARGS) +{ + ReturnSetInfo *rsinfo = (ReturnSetInfo *) fcinfo->resultinfo; + Tuplestorestate *tupstore; + TupleDesc tupdesc; + MemoryContext per_query_ctx; + MemoryContext oldcontext; + simple_cache_state *state = get_simple_cache_state(); + dht_iterator iterator; + simple_cache_entry *entry; + + if (rsinfo == NULL || !IsA(rsinfo, ReturnSetInfo)) + ereport(ERROR, + (errcode(ERRCODE_FEATURE_NOT_SUPPORTED), + errmsg("set-valued function called in context that cannot accept a set"))) + ; + if (!(rsinfo->allowedModes & SFRM_Materialize)) + ereport(ERROR, + (errcode(ERRCODE_FEATURE_NOT_SUPPORTED), + errmsg("materialize mode required, but it is not " \ + "allowed in this context"))); + + per_query_ctx = rsinfo->econtext->ecxt_per_query_memory; + oldcontext = MemoryContextSwitchTo(per_query_ctx); + + if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE) + elog(ERROR, "return type must be a row type"); + + tupstore = tuplestore_begin_heap(true, false, 1024 * 1024); + rsinfo->returnMode = SFRM_Materialize; + rsinfo->setResult = tupstore; + rsinfo->setDesc = tupdesc; + + MemoryContextSwitchTo(oldcontext); + + dht_iterate_begin(state->hash_table, &iterator, false); + while ((entry = dht_iterate_next(&iterator))) + { + Datum values[2]; + bool nulls[2] = {false, false}; + + Assert(tupdesc->natts == 2); + values[0] = CStringGetDatum(cstring_to_text(entry->key)); + values[1] = + CStringGetDatum(cstring_to_text_with_len(dsa_get_address(state->area, + entry->value), + entry->value_size)); + tuplestore_putvalues(tupstore, tupdesc, values, nulls); + } + dht_iterate_end(&iterator); + tuplestore_donestoring(tupstore); + + return (Datum) 0; +} + +Datum +simple_cache_set_size_limit(PG_FUNCTION_ARGS) +{ + simple_cache_state *state = get_simple_cache_state(); + int limit = PG_GETARG_INT32(0); + + dsa_set_size_limit(state->area, limit); + + return (Datum) 0; +} + +Datum +simple_cache_dump_area(PG_FUNCTION_ARGS) +{ + simple_cache_state *state = get_simple_cache_state(); + + dsa_dump(state->area); + + return (Datum) 0; +} + +Datum +simple_cache_dump_hash_table(PG_FUNCTION_ARGS) +{ + simple_cache_state *state = get_simple_cache_state(); + + dht_dump(state->hash_table); + + return (Datum) 0; +} diff --git a/src/test/modules/simple_cache/simple_cache.control b/src/test/modules/simple_cache/simple_cache.control new file mode 100644 index 0000000..bcddbd1 --- /dev/null +++ b/src/test/modules/simple_cache/simple_cache.control @@ -0,0 +1,5 @@ +# simple_cache extension +comment = 'Simple cache extension' +default_version = '1.0' +module_pathname = '$libdir/simple_cache' +relocatable = true \ No newline at end of file