From b0a60e9df9856d28b9730ce0d5f24e90de63e4bc Mon Sep 17 00:00:00 2001 From: Andrew Pogrebnoy Date: Thu, 11 Apr 2024 14:25:36 +0300 Subject: [PATCH] Use shmem for the encrypt buffer --- src/access/pg_tde_xlog.c | 68 +++++++++++++++++++++----------- src/include/access/pg_tde_xlog.h | 9 +++++ src/include/pg_tde_defines.h | 1 + src/pg_tde.c | 6 ++- 4 files changed, 60 insertions(+), 24 deletions(-) diff --git a/src/access/pg_tde_xlog.c b/src/access/pg_tde_xlog.c index d20d96a6..0e9ea9bd 100644 --- a/src/access/pg_tde_xlog.c +++ b/src/access/pg_tde_xlog.c @@ -12,10 +12,12 @@ #include "postgres.h" +#include "pg_tde_defines.h" #include "access/xlog.h" #include "access/xlog_internal.h" #include "access/xloginsert.h" #include "storage/bufmgr.h" +#include "storage/shmem.h" #include "utils/memutils.h" #include "access/pg_tde_tdemap.h" @@ -24,8 +26,7 @@ #include "encryption/enc_tde.h" -/* Buffer for the XLog encryption */ -static char *TDEXLogEncryptBuf; +static char *TDEXLogEncryptBuf = NULL; static void SetXLogPageIVPrefix(TimeLineID tli, XLogRecPtr lsn, char* iv_prefix); static int XLOGChooseNumBuffers(void); @@ -123,6 +124,7 @@ pg_tde_rmgr_identify(uint8 info) * if segment is encrypted. * We could also encrypt Records while adding them to the XLog Buf but it'll be the slowest (?). */ + static int XLOGChooseNumBuffers(void) { @@ -136,30 +138,49 @@ XLOGChooseNumBuffers(void) return xbuffers; } -void -TDEInitXLogSmgr(void) +/* + * Defines the size of the XLog encryption buffer + */ +Size +TDEXLogEncryptBuffSize() { - int xbuffers; - - /* - * Alloc memory for encrypition buffer. I should fit XLog buffers (XLOG_BLCKSZ * wal_buffers). - * We can't (re)alloc this buf in pg_tde_xlog_seg_write() based on the write sezie as - * it's called in the CRIT section, hence no allocations are allowed. - * - * TODO: - * - alloc in the shmem to save memory - * - ? alloc smaller (config option) and write in chunks (slower)? - */ + int xbuffers; + xbuffers = (XLOGbuffers == -1) ? XLOGChooseNumBuffers() : XLOGbuffers; - TDEXLogEncryptBuf = (char *) MemoryContextAlloc(TopMemoryContext, (Size) XLOG_BLCKSZ * xbuffers); - memset(TDEXLogEncryptBuf, 0, (Size) XLOG_BLCKSZ * 512); + return (Size) XLOG_BLCKSZ * xbuffers; +} +/* + * Alloc memory for encrypition buffer. + * + * It should fit XLog buffers (XLOG_BLCKSZ * wal_buffers). We can't + * (re)alloc this buf in pg_tde_xlog_seg_write() based on the write size as + * it's called in the CRIT section, hence no allocations are allowed. + * + * Access to this buffer happens during XLogWrite() call which should + * be called with WALWriteLock held, hence no need in extra locks. + */ +void +TDEXLogShmemInit(void) +{ + bool foundBuf; + + TDEXLogEncryptBuf = (char *) + TYPEALIGN(PG_IO_ALIGN_SIZE, + ShmemInitStruct("TDE XLog Encrypt Buffer", + XLOG_TDE_ENC_BUFF_ALIGNED_SIZE, + &foundBuf)); +} + +void +TDEInitXLogSmgr(void) +{ SetXLogSmgr(&tde_xlog_smgr); } /* * TODO: proper key management - * where to store the ref to the master and internal key? + * where to store refs to the master and internal keys? */ static InternalKey XLogInternalKey = {.key = {0xD,}}; @@ -175,8 +196,9 @@ pg_tde_xlog_seg_write(int fd, const void *buf, size_t count, off_t offset) Assert((count % (Size) XLOG_BLCKSZ) == 0); - elog(DEBUG1, "==> pg_tde_xlog_seg_WRITE, pages: %d", count / (Size) XLOG_BLCKSZ); - +#ifdef TDE_XLOG_DEBUG + elog(DEBUG1, "Write to a WAL segment, pages amount: %d", count / (Size) XLOG_BLCKSZ); +#endif /* Encrypt pages */ for (page_off = 0; page_off < count; page_off += (Size) XLOG_BLCKSZ) { @@ -207,7 +229,9 @@ pg_tde_xlog_seg_read(int fd, void *buf, size_t count, off_t offset) RelKeyData key = {.internal_key = XLogInternalKey}; char *decrypt_buf = NULL; - elog(DEBUG1, "==> pg_tde_xlog_seg_READ, pages: %d", count / (Size) XLOG_BLCKSZ); +#ifdef TDE_XLOG_DEBUG + elog(DEBUG1, "Read from a WAL segment, pages amount: %d", count / (Size) XLOG_BLCKSZ); +#endif readsz = pg_pread(fd, buf, count, offset); @@ -241,8 +265,6 @@ pg_tde_xlog_seg_read(int fd, void *buf, size_t count, off_t offset) static void SetXLogPageIVPrefix(TimeLineID tli, XLogRecPtr lsn, char* iv_prefix) { - elog(DEBUG1, "==> XlogIV %u, %lu", tli, lsn); - iv_prefix[0] = (tli >> 24); iv_prefix[1] = ((tli >> 16) & 0xFF); iv_prefix[2] = ((tli >> 8) & 0xFF); diff --git a/src/include/access/pg_tde_xlog.h b/src/include/access/pg_tde_xlog.h index 14aba09d..80601241 100644 --- a/src/include/access/pg_tde_xlog.h +++ b/src/include/access/pg_tde_xlog.h @@ -10,6 +10,7 @@ #define PG_TDE_XLOG_H #include "postgres.h" +#include "access/xlog.h" #include "access/xlog_internal.h" #include "access/xlog_smgr.h" @@ -34,6 +35,14 @@ static const RmgrData pg_tde_rmgr = { .rm_identify = pg_tde_rmgr_identify }; +/* XLog encryption staff */ + +extern Size TDEXLogEncryptBuffSize(); + +#define XLOG_TDE_ENC_BUFF_ALIGNED_SIZE add_size(TDEXLogEncryptBuffSize(), PG_IO_ALIGN_SIZE) + +extern void TDEXLogShmemInit(void); + extern ssize_t pg_tde_xlog_seg_read(int fd, void *buf, size_t count, off_t offset); extern ssize_t pg_tde_xlog_seg_write(int fd, const void *buf, size_t count, off_t offset); diff --git a/src/include/pg_tde_defines.h b/src/include/pg_tde_defines.h index aaa49722..0c9847bb 100644 --- a/src/include/pg_tde_defines.h +++ b/src/include/pg_tde_defines.h @@ -22,6 +22,7 @@ //#define ENCRYPTION_DEBUG 1 //#define KEYRING_DEBUG 1 //#define TDE_FORK_DEBUG 1 +// #define TDE_XLOG_DEBUG 1 #define pg_tde_fill_tuple heap_fill_tuple #define pg_tde_form_tuple heap_form_tuple diff --git a/src/pg_tde.c b/src/pg_tde.c index 324f7281..6070e0c9 100644 --- a/src/pg_tde.c +++ b/src/pg_tde.c @@ -59,6 +59,9 @@ tde_shmem_request(void) { Size sz = TdeRequiredSharedMemorySize(); int required_locks = TdeRequiredLocksCount(); + + sz = add_size(sz, XLOG_TDE_ENC_BUFF_ALIGNED_SIZE); + if (prev_shmem_request_hook) prev_shmem_request_hook(); RequestAddinShmemSpace(sz); @@ -74,6 +77,8 @@ tde_shmem_startup(void) TdeShmemInit(); AesInit(); + TDEXLogShmemInit(); + TDEInitXLogSmgr(); } void @@ -98,7 +103,6 @@ _PG_init(void) InstallFileKeyring(); InstallVaultV2Keyring(); RegisterCustomRmgr(RM_TDERMGR_ID, &pg_tde_rmgr); - TDEInitXLogSmgr(); } Datum pg_tde_extension_initialize(PG_FUNCTION_ARGS)