Add XLOG functionality to key provider info file and unify XLOG messa…

…ges for extension installation cleanup (#236)

* Add XLOG functionality to key provider info file and unify XLOG messages for extension installation cleanup

src/access/pg_tde_xlog.c

@@ -12,11 +12,13 @@
 
 #include "postgres.h"
 
+#include "pg_tde.h"
 #include "pg_tde_defines.h"
 #include "access/xlog.h"
 #include "access/xlog_internal.h"
 #include "access/xloginsert.h"
 #include "catalog/pg_tablespace_d.h"
+#include "catalog/tde_keyring.h"
 #include "storage/bufmgr.h"
 #include "storage/shmem.h"
 #include "utils/guc.h"
@@ -61,12 +63,19 @@ pg_tde_rmgr_redo(XLogReaderState *record)
 
 		save_principal_key_info(mkey);
 	}
-	else if (info == XLOG_TDE_CLEAN_PRINCIPAL_KEY)
+	else if (info == XLOG_TDE_EXTENSION_INSTALL_KEY)
 	{
-		XLogPrincipalKeyCleanup *xlrec = (XLogPrincipalKeyCleanup *) XLogRecGetData(record);
+		XLogExtensionInstall *xlrec = (XLogExtensionInstall *)XLogRecGetData(record);
 
-		cleanup_principal_key_info(xlrec->databaseId, xlrec->tablespaceId);
+		extension_install_redo(xlrec);
 	}
+
+	else if (info == XLOG_TDE_ADD_KEY_PROVIDER_KEY)
+	{
+		KeyringProviderXLRecord *xlrec = (KeyringProviderXLRecord *)XLogRecGetData(record);
+		redo_key_provider_info(xlrec);
+	}
+
 	else if (info == XLOG_TDE_ROTATE_KEY)
 	{
 		XLogPrincipalKeyRotate *xlrec = (XLogPrincipalKeyRotate *) XLogRecGetData(record);
@@ -96,18 +105,24 @@ pg_tde_rmgr_desc(StringInfo buf, XLogReaderState *record)
 
 		appendStringInfo(buf, "add tde principal key for db %u/%u", xlrec->databaseId, xlrec->tablespaceId);
 	}
-	if (info == XLOG_TDE_CLEAN_PRINCIPAL_KEY)
+	if (info == XLOG_TDE_EXTENSION_INSTALL_KEY)
 	{
-		XLogPrincipalKeyCleanup *xlrec = (XLogPrincipalKeyCleanup *) XLogRecGetData(record);
+		XLogExtensionInstall *xlrec = (XLogExtensionInstall *)XLogRecGetData(record);
 
-		appendStringInfo(buf, "cleanup tde principal key info for db %u/%u", xlrec->databaseId, xlrec->tablespaceId);
+		appendStringInfo(buf, "tde extension install for db %u/%u", xlrec->database_id, xlrec->tablespace_id);
 	}
 	if (info == XLOG_TDE_ROTATE_KEY)
 	{
 		XLogPrincipalKeyRotate *xlrec = (XLogPrincipalKeyRotate *) XLogRecGetData(record);
 
 		appendStringInfo(buf, "rotate principal key for %u", xlrec->databaseId);
 	}
+	if (info == XLOG_TDE_ADD_KEY_PROVIDER_KEY)
+	{
+		KeyringProviderXLRecord *xlrec = (KeyringProviderXLRecord *)XLogRecGetData(record);
+
+		appendStringInfo(buf, "add key provider %s for %u", xlrec->provider.provider_name, xlrec->database_id);
+	}
 }
 
 const char *
@@ -119,8 +134,8 @@ pg_tde_rmgr_identify(uint8 info)
 	if ((info & ~XLR_INFO_MASK) == XLOG_TDE_ADD_PRINCIPAL_KEY)
 		return "XLOG_TDE_ADD_PRINCIPAL_KEY";
 
-	if ((info & ~XLR_INFO_MASK) == XLOG_TDE_CLEAN_PRINCIPAL_KEY)
-		return "XLOG_TDE_CLEAN_PRINCIPAL_KEY";
+	if ((info & ~XLR_INFO_MASK) == XLOG_TDE_EXTENSION_INSTALL_KEY)
+		return "XLOG_TDE_EXTENSION_INSTALL_KEY";
 
 	return NULL;
 }

src/catalog/tde_keyring.c

@@ -9,7 +9,10 @@
  *
  *-------------------------------------------------------------------------
  */
-
+#include "postgres.h"
+#include "access/xlog.h"
+#include "access/xloginsert.h"
+#include "access/pg_tde_xlog.h"
 #include "catalog/tde_keyring.h"
 #include "catalog/tde_principal_key.h"
 #include "access/skey.h"
@@ -64,8 +67,9 @@ static GenericKeyring *load_keyring_provider_options(ProviderType provider_type,
 static VaultV2Keyring *load_vaultV2_keyring_provider_options(Datum keyring_options);
 static void debug_print_kerying(GenericKeyring *keyring);
 static char *get_keyring_infofile_path(char *resPath, Oid dbOid, Oid spcOid);
-static uint32 save_key_provider(KeyringProvideRecord *provider);
-static void key_provider_startup_cleanup(int tde_tbl_count, void *arg);
+static void key_provider_startup_cleanup(int tde_tbl_count, XLogExtensionInstall *ext_info, bool redo, void *arg);
+static uint32 write_key_provider_info(KeyringProvideRecord *provider, Oid database_id, Oid tablespace_id, off_t position, bool redo);
+static uint32 save_new_key_provider_info(KeyringProvideRecord *provider);
 
 static Size initialize_shared_state(void *start_address);
 static Size required_shared_mem_size(void);
@@ -112,7 +116,7 @@ void InitializeKeyProviderInfo(void)
 	on_ext_install(key_provider_startup_cleanup, NULL);
 }
 static void
-key_provider_startup_cleanup(int tde_tbl_count, void *arg)
+key_provider_startup_cleanup(int tde_tbl_count, XLogExtensionInstall *ext_info, bool redo, void *arg)
 {
 
 	if (tde_tbl_count > 0)
@@ -121,15 +125,7 @@ key_provider_startup_cleanup(int tde_tbl_count, void *arg)
 				(errmsg("failed to perform initialization. database already has %d TDE tables", tde_tbl_count)));
 		return;
 	}
-	cleanup_key_provider_info(MyDatabaseId, MyDatabaseTableSpace);
-
-	/* TODO: XLog the key cleanup */
-	// XLogPrincipalKeyCleanup xlrec;
-	// xlrec.databaseId = MyDatabaseId;
-	// xlrec.tablespaceId = MyDatabaseTableSpace;
-	// XLogBeginInsert();
-	// XLogRegisterData((char *)&xlrec, sizeof(TDEPrincipalKeyInfo));
-	// XLogInsert(RM_TDERMGR_ID, XLOG_TDE_CLEAN_PRINCIPAL_KEY);
+	cleanup_key_provider_info(ext_info->database_id, ext_info->tablespace_id);
 }
 
 ProviderType
@@ -314,11 +310,8 @@ fetch_next_key_provider(int fd, off_t* curr_pos, KeyringProvideRecord *provider)
 	return true;
 }
 
-/*
-* Save the key provider info to the file
-*/
 static uint32
-save_key_provider(KeyringProvideRecord *provider)
+write_key_provider_info(KeyringProvideRecord *provider, Oid database_id, Oid tablespace_id, off_t position, bool redo)
 {
 	off_t bytes_written = 0;
 	off_t curr_pos = 0;
@@ -329,7 +322,7 @@ save_key_provider(KeyringProvideRecord *provider)
 
 	Assert(provider != NULL);
 
-	get_keyring_infofile_path(kp_info_path, MyDatabaseId, MyDatabaseTableSpace);
+	get_keyring_infofile_path(kp_info_path, database_id, tablespace_id);
 
 	LWLockAcquire(tde_provider_info_lock(), LW_EXCLUSIVE);
 
@@ -341,44 +334,63 @@ save_key_provider(KeyringProvideRecord *provider)
 			(errcode_for_file_access(),
 				errmsg("could not open tde file \"%s\": %m", kp_info_path)));
 	}
-
-	/* we also need to verify the name conflict and generate the next provider ID */
-	while (fetch_next_key_provider(fd, &curr_pos, &existing_provider))
+	if (!redo)
 	{
-		if (strcmp(existing_provider.provider_name, provider->provider_name) == 0)
+		KeyringProviderXLRecord xlrec;
+		/* we also need to verify the name conflict and generate the next provider ID */
+		while (fetch_next_key_provider(fd, &curr_pos, &existing_provider))
 		{
-			close(fd);
-			LWLockRelease(tde_provider_info_lock());
-			ereport(ERROR,
-				(errcode(ERRCODE_DUPLICATE_OBJECT),
-					errmsg("key provider \"%s\" already exists", provider->provider_name)));
+			if (strcmp(existing_provider.provider_name, provider->provider_name) == 0)
+			{
+				close(fd);
+				LWLockRelease(tde_provider_info_lock());
+				ereport(ERROR,
+						(errcode(ERRCODE_DUPLICATE_OBJECT),
+						errmsg("key provider \"%s\" already exists", provider->provider_name)));
+			}
+			if (max_provider_id < existing_provider.provider_id)
+				max_provider_id = existing_provider.provider_id;
 		}
-		if (max_provider_id < existing_provider.provider_id)
-			max_provider_id = existing_provider.provider_id;
+		provider->provider_id = max_provider_id + 1;
+		curr_pos = lseek(fd, 0, SEEK_END);
+		/* emit the xlog here. So that we can handle partial file write errors */
+		xlrec.database_id = database_id;
+		xlrec.tablespace_id = tablespace_id;
+		xlrec.offset_in_file = curr_pos;
+		memcpy(&xlrec.provider, provider, sizeof(KeyringProvideRecord));
+
+		XLogBeginInsert();
+		XLogRegisterData((char *)&xlrec, sizeof(KeyringProviderXLRecord));
+		XLogInsert(RM_TDERMGR_ID, XLOG_TDE_ADD_KEY_PROVIDER_KEY);
+	}
+	else
+	{
+		/* we are performing redo, just go to the position received from the
+		 * xlog and write the record there.
+		 * No need to verify the name conflict and generate the provider ID
+		 */
+		curr_pos = lseek(fd, position, SEEK_SET);
 	}
-	provider->provider_id = max_provider_id + 1;
 	/*
 	 * All good, Just add a new provider
-	 * Write key to the end of file
 	 */
-	curr_pos = lseek(fd, 0, SEEK_END);
 	bytes_written = pg_pwrite(fd, provider, sizeof(KeyringProvideRecord), curr_pos);
 	if (bytes_written != sizeof(KeyringProvideRecord))
 	{
 		close(fd);
 		LWLockRelease(tde_provider_info_lock());
 		ereport(ERROR,
-			(errcode_for_file_access(),
-				errmsg("key provider info file \"%s\" can't be written: %m",
+				(errcode_for_file_access(),
+				 errmsg("key provider info file \"%s\" can't be written: %m",
 						kp_info_path)));
 	}
 	if (pg_fsync(fd) != 0)
 	{
 		close(fd);
 		LWLockRelease(tde_provider_info_lock());
 		ereport(ERROR,
-			(errcode_for_file_access(),
-				errmsg("could not fsync file \"%s\": %m",
+				(errcode_for_file_access(),
+				 errmsg("could not fsync file \"%s\": %m",
 						kp_info_path)));
 	}
 	close(fd);
@@ -387,10 +399,24 @@ save_key_provider(KeyringProvideRecord *provider)
 }
 
 /*
- * Scan the key provider info file and can also apply filter based on scanType
+ * Save the key provider info to the file
  */
-static List *
-scan_key_provider_file(ProviderScanType scanType, void* scanKey)
+static uint32
+save_new_key_provider_info(KeyringProvideRecord* provider)
+{
+	return write_key_provider_info(provider, MyDatabaseId, MyDatabaseTableSpace, 0, false);
+}
+
+uint32
+redo_key_provider_info(KeyringProviderXLRecord* xlrec)
+{
+	return write_key_provider_info(&xlrec->provider, xlrec->database_id, xlrec->tablespace_id, xlrec->offset_in_file, true);
+}
+
+/*
+	* Scan the key provider info file and can also apply filter based on scanType
+	*/
+static List *scan_key_provider_file(ProviderScanType scanType, void *scanKey)
 {
 	off_t curr_pos = 0;
 	int fd;
@@ -483,6 +509,7 @@ pg_tde_add_key_provider_internal(PG_FUNCTION_ARGS)
 	strncpy(provider.options, options, sizeof(provider.options));
 	strncpy(provider.provider_name, provider_name, sizeof(provider.provider_name));
 	provider.provider_type = get_keyring_provider_from_typename(provider_type);
-	save_key_provider(&provider);
+	save_new_key_provider_info(&provider);
+
 	PG_RETURN_INT32(provider.provider_id);
 }

src/catalog/tde_principal_key.c

@@ -66,7 +66,7 @@ static void initialize_objects_in_dsa_area(dsa_area *dsa, void *raw_dsa_area);
 static Size cache_area_size(void);
 static Size required_shared_mem_size(void);
 static void shared_memory_shutdown(int code, Datum arg);
-static void principal_key_startup_cleanup(int tde_tbl_count, void *arg);
+static void principal_key_startup_cleanup(int tde_tbl_count, XLogExtensionInstall *ext_info, bool redo, void *arg);
 static void clear_principal_key_cache(Oid databaseId) ;
 static inline dshash_table *get_principal_key_Hash(void);
 static TDEPrincipalKey *get_principal_key_from_cache(Oid dbOid);
@@ -658,25 +658,16 @@ push_principal_key_to_cache(TDEPrincipalKey *principalKey)
  * but unfortunately we do not have any such mechanism in PG.
 */
 static void
-principal_key_startup_cleanup(int tde_tbl_count, void* arg)
+principal_key_startup_cleanup(int tde_tbl_count, XLogExtensionInstall *ext_info, bool redo, void *arg)
 {
-    XLogPrincipalKeyCleanup xlrec;
-
     if (tde_tbl_count > 0)
     {
         ereport(WARNING,
                 (errmsg("Failed to perform initialization. database already has %d TDE tables", tde_tbl_count)));
         return;
     }
 
-    cleanup_principal_key_info(MyDatabaseId, MyDatabaseTableSpace);
-
-    /* XLog the key cleanup */
-    xlrec.databaseId = MyDatabaseId;
-    xlrec.tablespaceId = MyDatabaseTableSpace;
-    XLogBeginInsert();
-    XLogRegisterData((char *) &xlrec, sizeof(TDEPrincipalKeyInfo));
-    XLogInsert(RM_TDERMGR_ID, XLOG_TDE_CLEAN_PRINCIPAL_KEY);
+    cleanup_principal_key_info(ext_info->database_id, ext_info->tablespace_id);
 }
 
 void

src/include/access/pg_tde_xlog.h

@@ -17,10 +17,11 @@
 #endif
 
 /* TDE XLOG resource manager */
-#define XLOG_TDE_ADD_RELATION_KEY	0x00
+#define XLOG_TDE_ADD_RELATION_KEY		0x00
 #define XLOG_TDE_ADD_PRINCIPAL_KEY		0x10
-#define XLOG_TDE_CLEAN_PRINCIPAL_KEY	0x20
-#define XLOG_TDE_ROTATE_KEY			0x30
+#define XLOG_TDE_EXTENSION_INSTALL_KEY	0x20
+#define XLOG_TDE_ROTATE_KEY				0x30
+#define XLOG_TDE_ADD_KEY_PROVIDER_KEY 	0x40
 
 /* TODO: ID has to be registedred and changed: https://wiki.postgresql.org/wiki/CustomWALResourceManagers */
 #define RM_TDERMGR_ID	RM_EXPERIMENTAL_ID

src/include/catalog/tde_keyring.h

@@ -69,11 +69,19 @@ typedef struct KeyringProvideRecord
 	char options[MAX_KEYRING_OPTION_LEN];
 	ProviderType provider_type;
 } KeyringProvideRecord;
+typedef struct KeyringProviderXLRecord
+{
+	Oid database_id;
+	Oid tablespace_id;
+	off_t offset_in_file;
+	KeyringProvideRecord provider;
+} KeyringProviderXLRecord;
 
 extern List *GetAllKeyringProviders(void);
 extern GenericKeyring *GetKeyProviderByName(const char *provider_name);
 extern GenericKeyring *GetKeyProviderByID(int provider_id);
 extern ProviderType get_keyring_provider_from_typename(char *provider_type);
 extern void cleanup_key_provider_info(Oid databaseId, Oid tablespaceId);
 extern void InitializeKeyProviderInfo(void);
+extern uint32 redo_key_provider_info(KeyringProviderXLRecord *xlrec);
 #endif /*TDE_KEYRING_H*/

src/include/catalog/tde_principal_key.h

@@ -55,12 +55,6 @@ typedef struct XLogPrincipalKeyRotate
 
 #define SizeoOfXLogPrincipalKeyRotate	offsetof(XLogPrincipalKeyRotate, buff)
 
-typedef struct XLogPrincipalKeyCleanup
-{
-	Oid databaseId;
-	Oid tablespaceId;
-} XLogPrincipalKeyCleanup;
-
 extern void InitializePrincipalKeyInfo(void);
 extern void cleanup_principal_key_info(Oid databaseId, Oid tablespaceId);
 extern LWLock *tde_lwlock_mk_files(void);

src/include/pg_tde.h

@@ -8,8 +8,15 @@
 #ifndef PG_TDE_H
 #define PG_TDE_H
 
-typedef void (*pg_tde_on_ext_install_callback)(int tde_tbl_count, void* arg);
+typedef struct XLogExtensionInstall
+{
+    Oid database_id;
+    Oid tablespace_id;
+} XLogExtensionInstall;
+
+typedef void (*pg_tde_on_ext_install_callback)(int tde_tbl_count, XLogExtensionInstall* ext_info, bool redo, void *arg);
 
 extern void on_ext_install(pg_tde_on_ext_install_callback function, void* arg);
 
+extern void extension_install_redo(XLogExtensionInstall *xlrec);
 #endif /*PG_TDE_H*/