6666#include "postgres.h"
6767
6868#include "access/transam.h"
69+ #include "fmgr.h"
70+ #include "funcapi.h"
6971#include "miscadmin.h"
7072#include "storage/lwlock.h"
7173#include "storage/pg_shmem.h"
7274#include "storage/shmem.h"
7375#include "storage/spin.h"
76+ #include "utils/builtins.h"
7477
78+ static void * ShmemAllocRaw (Size size , Size * allocated_size );
7579
7680/* shared memory global variables */
7781
@@ -157,8 +161,9 @@ void *
157161ShmemAlloc (Size size )
158162{
159163 void * newSpace ;
164+ Size allocated_size ;
160165
161- newSpace = ShmemAllocNoError (size );
166+ newSpace = ShmemAllocRaw (size , & allocated_size );
162167 if (!newSpace )
163168 ereport (ERROR ,
164169 (errcode (ERRCODE_OUT_OF_MEMORY ),
@@ -174,6 +179,20 @@ ShmemAlloc(Size size)
174179 */
175180void *
176181ShmemAllocNoError (Size size )
182+ {
183+ Size allocated_size ;
184+
185+ return ShmemAllocRaw (size , & allocated_size );
186+ }
187+
188+ /*
189+ * ShmemAllocRaw -- allocate align chunk and return allocated size
190+ *
191+ * Also sets *allocated_size to the number of bytes allocated, which will
192+ * be equal to the number requested plus any padding we choose to add.
193+ */
194+ static void *
195+ ShmemAllocRaw (Size size , Size * allocated_size )
177196{
178197 Size newStart ;
179198 Size newFree ;
@@ -191,6 +210,7 @@ ShmemAllocNoError(Size size)
191210 * won't be sufficient.
192211 */
193212 size = CACHELINEALIGN (size );
213+ * allocated_size = size ;
194214
195215 Assert (ShmemSegHdr != NULL );
196216
@@ -441,8 +461,10 @@ ShmemInitStruct(const char *name, Size size, bool *foundPtr)
441461 }
442462 else
443463 {
464+ Size allocated_size ;
465+
444466 /* It isn't in the table yet. allocate and initialize it */
445- structPtr = ShmemAllocNoError (size );
467+ structPtr = ShmemAllocRaw (size , & allocated_size );
446468 if (structPtr == NULL )
447469 {
448470 /* out of memory; remove the failed ShmemIndex entry */
@@ -455,6 +477,7 @@ ShmemInitStruct(const char *name, Size size, bool *foundPtr)
455477 name , size )));
456478 }
457479 result -> size = size ;
480+ result -> allocated_size = allocated_size ;
458481 result -> location = structPtr ;
459482 }
460483
@@ -503,3 +526,82 @@ mul_size(Size s1, Size s2)
503526 errmsg ("requested shared memory size overflows size_t" )));
504527 return result ;
505528}
529+
530+ /* SQL SRF showing allocated shared memory */
531+ Datum
532+ pg_get_shmem_allocations (PG_FUNCTION_ARGS )
533+ {
534+ #define PG_GET_SHMEM_SIZES_COLS 4
535+ ReturnSetInfo * rsinfo = (ReturnSetInfo * ) fcinfo -> resultinfo ;
536+ TupleDesc tupdesc ;
537+ Tuplestorestate * tupstore ;
538+ MemoryContext per_query_ctx ;
539+ MemoryContext oldcontext ;
540+ HASH_SEQ_STATUS hstat ;
541+ ShmemIndexEnt * ent ;
542+ Size named_allocated = 0 ;
543+ Datum values [PG_GET_SHMEM_SIZES_COLS ];
544+ bool nulls [PG_GET_SHMEM_SIZES_COLS ];
545+
546+ /* check to see if caller supports us returning a tuplestore */
547+ if (rsinfo == NULL || !IsA (rsinfo , ReturnSetInfo ))
548+ ereport (ERROR ,
549+ (errcode (ERRCODE_FEATURE_NOT_SUPPORTED ),
550+ errmsg ("set-valued function called in context that cannot accept a set" )));
551+ if (!(rsinfo -> allowedModes & SFRM_Materialize ))
552+ ereport (ERROR ,
553+ (errcode (ERRCODE_FEATURE_NOT_SUPPORTED ),
554+ errmsg ("materialize mode required, but it is not allowed in this context" )));
555+
556+ /* Build a tuple descriptor for our result type */
557+ if (get_call_result_type (fcinfo , NULL , & tupdesc ) != TYPEFUNC_COMPOSITE )
558+ elog (ERROR , "return type must be a row type" );
559+
560+ per_query_ctx = rsinfo -> econtext -> ecxt_per_query_memory ;
561+ oldcontext = MemoryContextSwitchTo (per_query_ctx );
562+
563+ tupstore = tuplestore_begin_heap (true, false, work_mem );
564+ rsinfo -> returnMode = SFRM_Materialize ;
565+ rsinfo -> setResult = tupstore ;
566+ rsinfo -> setDesc = tupdesc ;
567+
568+ MemoryContextSwitchTo (oldcontext );
569+
570+ LWLockAcquire (ShmemIndexLock , LW_SHARED );
571+
572+ hash_seq_init (& hstat , ShmemIndex );
573+
574+ /* output all allocated entries */
575+ memset (nulls , 0 , sizeof (nulls ));
576+ while ((ent = (ShmemIndexEnt * ) hash_seq_search (& hstat )) != NULL )
577+ {
578+ values [0 ] = CStringGetTextDatum (ent -> key );
579+ values [1 ] = Int64GetDatum ((char * ) ent -> location - (char * ) ShmemSegHdr );
580+ values [2 ] = Int64GetDatum (ent -> size );
581+ values [3 ] = Int64GetDatum (ent -> allocated_size );
582+ named_allocated += ent -> allocated_size ;
583+
584+ tuplestore_putvalues (tupstore , tupdesc , values , nulls );
585+ }
586+
587+ /* output shared memory allocated but not counted via the shmem index */
588+ values [0 ] = CStringGetTextDatum ("<anonymous>" );
589+ nulls [1 ] = true;
590+ values [2 ] = Int64GetDatum (ShmemSegHdr -> freeoffset - named_allocated );
591+ values [3 ] = values [2 ];
592+ tuplestore_putvalues (tupstore , tupdesc , values , nulls );
593+
594+ /* output as-of-yet unused shared memory */
595+ nulls [0 ] = true;
596+ values [1 ] = Int64GetDatum (ShmemSegHdr -> freeoffset );
597+ nulls [1 ] = false;
598+ values [2 ] = Int64GetDatum (ShmemSegHdr -> totalsize - ShmemSegHdr -> freeoffset );
599+ values [3 ] = values [2 ];
600+ tuplestore_putvalues (tupstore , tupdesc , values , nulls );
601+
602+ LWLockRelease (ShmemIndexLock );
603+
604+ tuplestore_donestoring (tupstore );
605+
606+ return (Datum ) 0 ;
607+ }
0 commit comments