use_work_mem_as_max_bytes.patch

diff --git a/src/test/modules/test_tidstore/expected/test_tidstore.out b/src/test/modules/test_tidstore/expected/test_tidstore.out
index cbcacfd26e..07be72683a 100644
--- a/src/test/modules/test_tidstore/expected/test_tidstore.out
+++ b/src/test/modules/test_tidstore/expected/test_tidstore.out
@@ -6,6 +6,8 @@ CREATE TEMP TABLE hideblocks(blockno bigint);
 -- The maximum number of heap tuples (MaxHeapTuplesPerPage) in 8kB block is 291.
 -- We use a higher number to test tidstore.
 \set maxoffset 512
+-- We use work_mem as the maximum bytes in the tests.
+SET work_mem To '2MB';
 SELECT test_create(false);
  test_create 
 -------------
@@ -230,4 +232,33 @@ SELECT test_destroy();
  
 (1 row)
 
+-- Test to create tidstore with a small maximum byte, using small memory block
+-- sizes internally.
+SET work_mem TO '64kB';
+SELECT test_create(false);
+ test_create 
+-------------
+ 
+(1 row)
+
+SELECT test_destroy();
+ test_destroy 
+--------------
+ 
+(1 row)
+
+-- Same test for shared tidstore with 1MB, the minimum value of maintenance_work_mem.
+SET work_mem TO '1MB';
+SELECT test_create(true);
+ test_create 
+-------------
+ 
+(1 row)
+
+SELECT test_destroy();
+ test_destroy 
+--------------
+ 
+(1 row)
+
 DROP TABLE hideblocks;
diff --git a/src/test/modules/test_tidstore/sql/test_tidstore.sql b/src/test/modules/test_tidstore/sql/test_tidstore.sql
index a29e4ec1c5..71cb6f20b0 100644
--- a/src/test/modules/test_tidstore/sql/test_tidstore.sql
+++ b/src/test/modules/test_tidstore/sql/test_tidstore.sql
@@ -9,6 +9,8 @@ CREATE TEMP TABLE hideblocks(blockno bigint);
 -- We use a higher number to test tidstore.
 \set maxoffset 512
 
+-- We use work_mem as the maximum bytes in the tests.
+SET work_mem To '2MB';
 SELECT test_create(false);
 
 -- Test on empty tidstore.
@@ -83,4 +85,16 @@ SELECT check_set_block_offsets();
 
 -- cleanup
 SELECT test_destroy();
+
+-- Test to create tidstore with a small maximum byte, using small memory block
+-- sizes internally.
+SET work_mem TO '64kB';
+SELECT test_create(false);
+SELECT test_destroy();
+
+-- Same test for shared tidstore with 1MB, the minimum value of maintenance_work_mem.
+SET work_mem TO '1MB';
+SELECT test_create(true);
+SELECT test_destroy();
+
 DROP TABLE hideblocks;
diff --git a/src/test/modules/test_tidstore/test_tidstore.c b/src/test/modules/test_tidstore/test_tidstore.c
index 5417163407..af00f0ce69 100644
--- a/src/test/modules/test_tidstore/test_tidstore.c
+++ b/src/test/modules/test_tidstore/test_tidstore.c
@@ -19,6 +19,7 @@
 #include "access/tidstore.h"
 #include "fmgr.h"
 #include "funcapi.h"
+#include "miscadmin.h"
 #include "storage/block.h"
 #include "storage/itemptr.h"
 #include "storage/lwlock.h"
@@ -89,7 +90,7 @@ test_create(PG_FUNCTION_ARGS)
 	MemoryContext old_ctx;
 
 	/* doesn't really matter, since it's just a hint */
-	size_t		tidstore_max_size = 2 * 1024 * 1024;
+	size_t		tidstore_max_size = work_mem * 1024;
 	size_t		array_init_size = 1024;
 
 	Assert(tidstore == NULL);