Skip to content

PG-1604 fix: preallocate one more record for the cache #539

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Aug 18, 2025
Merged

PG-1604 fix: preallocate one more record for the cache #539

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 36 additions & 2 deletions contrib/pg_tde/src/access/pg_tde_xlog_keys.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
#include "common/pg_tde_utils.h"
#include "encryption/enc_aes.h"
#include "encryption/enc_tde.h"
#include "utils/palloc.h"

#ifdef FRONTEND
#include "pg_tde_fe.h"
Expand Down Expand Up @@ -43,7 +44,9 @@ typedef struct WalKeyFileEntry
} WalKeyFileEntry;

static WALKeyCacheRec *tde_wal_key_cache = NULL;
static WALKeyCacheRec *tde_wal_prealloc_record = NULL;
static WALKeyCacheRec *tde_wal_key_last_rec = NULL;
static WalEncryptionKey *tde_wal_prealloc_key = NULL;

static WALKeyCacheRec *pg_tde_add_wal_key_to_cache(WalEncryptionKey *cached_key);
static WalEncryptionKey *pg_tde_decrypt_wal_key(const TDEPrincipalKey *principal_key, WalKeyFileEntry *entry);
Expand Down Expand Up @@ -326,6 +329,34 @@ pg_tde_fetch_wal_keys(WalLocation start)
return return_wal_rec;
}

/*
* In special cases, we have to add one more record to the WAL key cache during write (in the critical section, when we can't allocate).
* This method is a helper to deal with that: when adding a single key, we potentially allocate 2 records.
* These variables help preallocate them, so in the critical section we can just use the already allocated objects, and update the cache with them.
*
* While this is somewhat a hack, it is also simple, safe, reliable, and way easier to implement than to refactor the cache or the decryption/encryption loop.
*/
void
pg_tde_wal_cache_extra_palloc(void)
{
#ifndef FRONTEND
MemoryContext oldCtx;

oldCtx = MemoryContextSwitchTo(TopMemoryContext);
#endif
if (tde_wal_prealloc_record == NULL)
{
tde_wal_prealloc_record = palloc0_object(WALKeyCacheRec);
}
if (tde_wal_prealloc_key == NULL)
{
tde_wal_prealloc_key = palloc0_object(WalEncryptionKey);
}
#ifndef FRONTEND
MemoryContextSwitchTo(oldCtx);
#endif
}

static WALKeyCacheRec *
pg_tde_add_wal_key_to_cache(WalEncryptionKey *key)
{
Expand All @@ -335,7 +366,8 @@ pg_tde_add_wal_key_to_cache(WalEncryptionKey *key)

oldCtx = MemoryContextSwitchTo(TopMemoryContext);
#endif
wal_rec = palloc0_object(WALKeyCacheRec);
wal_rec = tde_wal_prealloc_record == NULL ? palloc0_object(WALKeyCacheRec) : tde_wal_prealloc_record;
tde_wal_prealloc_record = NULL;
#ifndef FRONTEND
MemoryContextSwitchTo(oldCtx);
#endif
Expand Down Expand Up @@ -553,7 +585,9 @@ pg_tde_write_wal_key_file_entry(const WalEncryptionKey *rel_key_data,
static WalEncryptionKey *
pg_tde_decrypt_wal_key(const TDEPrincipalKey *principal_key, WalKeyFileEntry *entry)
{
WalEncryptionKey *key = palloc_object(WalEncryptionKey);
WalEncryptionKey *key = tde_wal_prealloc_key == NULL ? palloc_object(WalEncryptionKey) : tde_wal_prealloc_key;

tde_wal_prealloc_key = NULL;

Assert(principal_key);

Expand Down
11 changes: 5 additions & 6 deletions contrib/pg_tde/src/access/pg_tde_xlog_smgr.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -248,6 +248,7 @@ TDEXLogSmgrInitWrite(bool encrypt_xlog)

/* cache is empty, prefetch keys from disk */
pg_tde_fetch_wal_keys(start);
pg_tde_wal_cache_extra_palloc();
}

if (key)
Expand Down Expand Up @@ -284,8 +285,8 @@ TDEXLogWriteEncryptedPagesOldKeys(int fd, const void *buf, size_t count, off_t o
memcpy(enc_buff, buf, count);

/*
* This method potentially allocates, but only in very early execution
* Shouldn't happen in a write, where we are in a critical section
* This method potentially allocates, but only in very early execution Can
* happen during a write, but we have one more cache entry preallocated.
*/
TDEXLogCryptBuffer(buf, enc_buff, count, offset, tli, segno, segSize);

Expand Down Expand Up @@ -356,7 +357,7 @@ tdeheap_xlog_seg_write(int fd, const void *buf, size_t count, off_t offset,
lastKeyUsable = (writeKeyLoc.lsn != 0);
afterWriteKey = wal_location_cmp(writeKeyLoc, loc) <= 0;

if (EncryptionKey.type != WAL_KEY_TYPE_INVALID && !lastKeyUsable)
if (EncryptionKey.type != WAL_KEY_TYPE_INVALID && !lastKeyUsable && afterWriteKey)
{
WALKeyCacheRec *last_key = pg_tde_get_last_wal_key();

Expand Down Expand Up @@ -442,10 +443,8 @@ TDEXLogCryptBuffer(const void *buf, void *out_buf, size_t count, off_t offset,
WALKeyCacheRec *last_key = pg_tde_get_last_wal_key();
WalLocation write_loc = {.tli = TDEXLogGetEncKeyTli(),.lsn = write_key_lsn};

Assert(last_key);

/* write has generated a new key, need to fetch it */
if (wal_location_cmp(last_key->start, write_loc) < 0)
if (last_key != NULL && wal_location_cmp(last_key->start, write_loc) < 0)
{
pg_tde_fetch_wal_keys(write_loc);

Expand Down
1 change: 1 addition & 0 deletions contrib/pg_tde/src/include/access/pg_tde_xlog_keys.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -82,5 +82,6 @@ extern WalEncryptionKey *pg_tde_read_last_wal_key(void);
extern void pg_tde_save_server_key(const TDEPrincipalKey *principal_key, bool write_xlog);
extern void pg_tde_save_server_key_redo(const TDESignedPrincipalKeyInfo *signed_key_info);
extern void pg_tde_wal_last_key_set_location(WalLocation loc);
extern void pg_tde_wal_cache_extra_palloc(void);

#endif /* PG_TDE_XLOG_KEYS_H */
Loading