Fix memory usage in the internal keys cache (#385)

There was a mistake in the new cap calculations during the cache extension. It popped up only when a new cache size was multiple of a cache record (every 256 records).  Which lead to the usage of the memory beyond an allocated size. This commit fixes it along with mlock size for reallocated pages.

Also fixed a typo in a variable name.

Fixes PG-1248

expected/cache_alloc.out

@@ -0,0 +1,125 @@
+-- We test cache so AM doesn't matter
+-- Just checking there are no mem debug WARNINGs during the cache population
+CREATE EXTENSION pg_tde;
+SELECT pg_tde_add_key_provider_file('file-vault','/tmp/pg_tde_test_keyring.per');
+ pg_tde_add_key_provider_file 
+------------------------------
+                            1
+(1 row)
+
+SELECT pg_tde_set_principal_key('test-db-principal-key','file-vault');
+ pg_tde_set_principal_key 
+--------------------------
+ t
+(1 row)
+
+do $$
+    DECLARE idx integer;
+begin
+    for idx in 0..700 loop
+        EXECUTE format('CREATE TABLE t%s (c1 int) USING tde_heap_basic', idx);
+    end loop;
+end; $$;
+DROP EXTENSION pg_tde cascade;
+NOTICE:  drop cascades to 701 other objects
+DETAIL:  drop cascades to table t0
+drop cascades to table t1
+drop cascades to table t2
+drop cascades to table t3
+drop cascades to table t4
+drop cascades to table t5
+drop cascades to table t6
+drop cascades to table t7
+drop cascades to table t8
+drop cascades to table t9
+drop cascades to table t10
+drop cascades to table t11
+drop cascades to table t12
+drop cascades to table t13
+drop cascades to table t14
+drop cascades to table t15
+drop cascades to table t16
+drop cascades to table t17
+drop cascades to table t18
+drop cascades to table t19
+drop cascades to table t20
+drop cascades to table t21
+drop cascades to table t22
+drop cascades to table t23
+drop cascades to table t24
+drop cascades to table t25
+drop cascades to table t26
+drop cascades to table t27
+drop cascades to table t28
+drop cascades to table t29
+drop cascades to table t30
+drop cascades to table t31
+drop cascades to table t32
+drop cascades to table t33
+drop cascades to table t34
+drop cascades to table t35
+drop cascades to table t36
+drop cascades to table t37
+drop cascades to table t38
+drop cascades to table t39
+drop cascades to table t40
+drop cascades to table t41
+drop cascades to table t42
+drop cascades to table t43
+drop cascades to table t44
+drop cascades to table t45
+drop cascades to table t46
+drop cascades to table t47
+drop cascades to table t48
+drop cascades to table t49
+drop cascades to table t50
+drop cascades to table t51
+drop cascades to table t52
+drop cascades to table t53
+drop cascades to table t54
+drop cascades to table t55
+drop cascades to table t56
+drop cascades to table t57
+drop cascades to table t58
+drop cascades to table t59
+drop cascades to table t60
+drop cascades to table t61
+drop cascades to table t62
+drop cascades to table t63
+drop cascades to table t64
+drop cascades to table t65
+drop cascades to table t66
+drop cascades to table t67
+drop cascades to table t68
+drop cascades to table t69
+drop cascades to table t70
+drop cascades to table t71
+drop cascades to table t72
+drop cascades to table t73
+drop cascades to table t74
+drop cascades to table t75
+drop cascades to table t76
+drop cascades to table t77
+drop cascades to table t78
+drop cascades to table t79
+drop cascades to table t80
+drop cascades to table t81
+drop cascades to table t82
+drop cascades to table t83
+drop cascades to table t84
+drop cascades to table t85
+drop cascades to table t86
+drop cascades to table t87
+drop cascades to table t88
+drop cascades to table t89
+drop cascades to table t90
+drop cascades to table t91
+drop cascades to table t92
+drop cascades to table t93
+drop cascades to table t94
+drop cascades to table t95
+drop cascades to table t96
+drop cascades to table t97
+drop cascades to table t98
+drop cascades to table t99
+and 601 other objects (see server log for list)

meson.build

@@ -104,6 +104,7 @@ sql_tests = [
       'kmip_test_basic',
       'alter_index_basic',
       'merge_join_basic',
+      'cache_alloc',
 ]
 
 tap_tests = [

sql/cache_alloc.sql

@@ -0,0 +1,17 @@
+-- We test cache so AM doesn't matter
+-- Just checking there are no mem debug WARNINGs during the cache population
+
+CREATE EXTENSION pg_tde;
+
+SELECT pg_tde_add_key_provider_file('file-vault','/tmp/pg_tde_test_keyring.per');
+SELECT pg_tde_set_principal_key('test-db-principal-key','file-vault');
+
+do $$
+    DECLARE idx integer;
+begin
+    for idx in 0..700 loop
+        EXECUTE format('CREATE TABLE t%s (c1 int) USING tde_heap_basic', idx);
+    end loop;
+end; $$;
+
+DROP EXTENSION pg_tde cascade;

src/access/pg_tde_tdemap.c

@@ -1460,40 +1460,49 @@ pg_tde_put_key_into_cache(RelFileNumber rel_num, RelKeyData *key)
 			elog(ERROR, "could not mlock internal key initial cache page: %m");
 
 		tde_rel_key_cache->len = 0;
-		tde_rel_key_cache->cap = pageSize / sizeof(RelKeyCacheRec);
+		tde_rel_key_cache->cap = (pageSize - 1) / sizeof(RelKeyCacheRec);
 	}
 
 	/*
 	 * Add another mem page if there is no more room left for another key. We
 	 * allocate `current_memory_size` + 1 page and copy data there.
 	 */
-	if (tde_rel_key_cache->len + 1 >
-		(tde_rel_key_cache->cap * sizeof(RelKeyCacheRec)) / sizeof(RelKeyCacheRec))
+	if (tde_rel_key_cache->len == tde_rel_key_cache->cap)
 	{
 		size_t size;
 		size_t old_size;
-		RelKeyCacheRec *chachePage;
+		RelKeyCacheRec *cachePage;
 
-		size = TYPEALIGN(pageSize, (tde_rel_key_cache->cap + 1) * sizeof(RelKeyCacheRec));
 		old_size = TYPEALIGN(pageSize, (tde_rel_key_cache->cap) * sizeof(RelKeyCacheRec));
 
+		/* TODO: consider some formula for less allocations when  caching a lot
+		 * of objects. But on the other, hand it'll use more memory...
+		 * E.g.:
+		 *	if (old_size < 0x8000)
+		 *		size = old_size * 2;
+		 *	else
+		 *		size = TYPEALIGN(pageSize, old_size + ((old_size + 3*256) >> 2));
+		 *		
+		 */
+		size = old_size + pageSize;
+
 #ifndef FRONTEND
 		oldCtx = MemoryContextSwitchTo(TopMemoryContext);
-		chachePage = palloc_aligned(size, pageSize, MCXT_ALLOC_ZERO);
+		cachePage = palloc_aligned(size, pageSize, MCXT_ALLOC_ZERO);
 		MemoryContextSwitchTo(oldCtx);
 #else
-		chachePage = aligned_alloc(pageSize, size);
-		memset(chachePage, 0, size);
+		cachePage = aligned_alloc(pageSize, size);
+		memset(cachePage, 0, size);
 #endif
 
-		memcpy(chachePage, tde_rel_key_cache->data, old_size);
+		memcpy(cachePage, tde_rel_key_cache->data, old_size);
 		pfree(tde_rel_key_cache->data);
-		tde_rel_key_cache->data = chachePage;
+		tde_rel_key_cache->data = cachePage;
 
-		if (mlock(tde_rel_key_cache->data, pageSize) == -1)
-			elog(ERROR, "could not mlock internal key cache page: %m");
+		if (mlock(tde_rel_key_cache->data, size) == -1)
+			elog(WARNING, "could not mlock internal key cache pages: %m");
 
-		tde_rel_key_cache->cap = size / sizeof(RelKeyCacheRec);
+		tde_rel_key_cache->cap = (size - 1) / sizeof(RelKeyCacheRec);
 	}
 
 	rec = tde_rel_key_cache->data + tde_rel_key_cache->len;