Add disk usage statistics

src/backend/distributed/metadata/metadata_utility.c

@@ -9,6 +9,7 @@
  *
  *-------------------------------------------------------------------------
  */
+#include <sys/statvfs.h>
 
 #include "postgres.h"
 #include "funcapi.h"
@@ -67,6 +68,9 @@
 
 
 /* Local functions forward declarations */
+static bool GetNodeDiskSpaceStats(char *nodeName, int nodePort, uint64 *availableBytes,
+								  uint64 *totalBytes);
+static bool GetLocalDiskSpaceStats(uint64 *availableBytes, uint64 *totalBytes);
 static uint64 NodeDatabaseSize(char *nodeName, int nodePort, char *databaseName);
 static uint64 * AllocateUint64(uint64 value);
 static void RecordDistributedRelationDependencies(Oid distributedRelationId);
@@ -82,6 +86,8 @@ static void ErrorIfNotSuitableToGetSize(Oid relationId);
 
 
 /* exports for SQL callable functions */
+PG_FUNCTION_INFO_V1(citus_node_disk_space_stats);
+PG_FUNCTION_INFO_V1(citus_local_disk_space_stats);
 PG_FUNCTION_INFO_V1(citus_database_size);
 PG_FUNCTION_INFO_V1(citus_node_database_size);
 PG_FUNCTION_INFO_V1(citus_table_size);
@@ -89,6 +95,158 @@ PG_FUNCTION_INFO_V1(citus_total_relation_size);
 PG_FUNCTION_INFO_V1(citus_relation_size);
 
 
+/*
+ * citus_node_database_size is a UDF that returns the size of the database
+ * on a given node.
+ */
+Datum
+citus_node_disk_space_stats(PG_FUNCTION_ARGS)
+{
+	text *nodeNameText = PG_GETARG_TEXT_P(0);
+	char *nodeNameString = text_to_cstring(nodeNameText);
+	int32 nodePort = PG_GETARG_INT32(1);
+
+	CheckCitusVersion(ERROR);
+
+	uint64 availableBytes = 0;
+	uint64 totalBytes = 0;
+
+	GetNodeDiskSpaceStats(nodeNameString, nodePort, &availableBytes, &totalBytes);
+
+	TupleDesc tupleDescriptor = NULL;
+	Datum values[TABLE_METADATA_FIELDS];
+	bool isNulls[TABLE_METADATA_FIELDS];
+
+	TypeFuncClass resultTypeClass = get_call_result_type(fcinfo, NULL,
+														 &tupleDescriptor);
+	if (resultTypeClass != TYPEFUNC_COMPOSITE)
+	{
+		ereport(ERROR, (errmsg("return type must be a row type")));
+	}
+
+	/* form heap tuple for remote disk space statistics */
+	memset(values, 0, sizeof(values));
+	memset(isNulls, false, sizeof(isNulls));
+
+	values[0] = UInt64GetDatum(availableBytes);
+	values[1] = UInt64GetDatum(totalBytes);
+
+	HeapTuple diskSpaceTuple = heap_form_tuple(tupleDescriptor, values, isNulls);
+
+	PG_RETURN_UINT64(HeapTupleGetDatum(diskSpaceTuple));
+}
+
+
+/*
+ * GetNodeDiskSpaceStats fetches the disk space statistics for a given node,
+ * or returns false if unsuccessful.
+ */
+static bool
+GetNodeDiskSpaceStats(char *nodeName, int nodePort, uint64 *availableBytes,
+					  uint64 *totalBytes)
+{
+	uint32 connectionFlags = 0;
+	PGresult *result = NULL;
+	bool raiseErrors = true;
+
+	char *sizeQuery = "SELECT available_disk_size, total_disk_size "
+					  "FROM pg_catalog.citus_local_disk_space_stats()";
+
+	MultiConnection *connection = GetNodeConnection(connectionFlags, nodeName, nodePort);
+
+	int queryResult = ExecuteOptionalRemoteCommand(connection, sizeQuery, &result);
+	if (queryResult != 0)
+	{
+		ereport(WARNING, (errcode(ERRCODE_CONNECTION_FAILURE),
+						  errmsg("cannot get the disk space statistics for node %s:%d",
+								 nodeName, nodePort)));
+		return false;
+	}
+
+	if (!IsResponseOK(result) || PQntuples(result) != 1)
+	{
+		ereport(WARNING, (errcode(ERRCODE_CONNECTION_FAILURE),
+						  errmsg("cannot get the disk space statistics for node %s:%d",
+								 nodeName, nodePort)));
+
+		PQclear(result);
+		ClearResults(connection, raiseErrors);
+
+		return false;
+	}
+
+	char *availableBytesString = PQgetvalue(result, 0, 0);
+	char *totalBytesString = PQgetvalue(result, 0, 1);
+
+	*availableBytes = SafeStringToUint64(availableBytesString);
+	*totalBytes = SafeStringToUint64(totalBytesString);
+
+	PQclear(result);
+	ClearResults(connection, raiseErrors);
+
+	return true;
+}
+
+
+/*
+ * citus_local_disk_space_stats returns total disk space and available disk
+ * space for the disk that contains PGDATA.
+ */
+Datum
+citus_local_disk_space_stats(PG_FUNCTION_ARGS)
+{
+	uint64 availableBytes = 0;
+	uint64 totalBytes = 0;
+
+	if (GetLocalDiskSpaceStats(&availableBytes, &totalBytes))
+	{
+		ereport(WARNING, (errmsg("could not get disk space")));
+	}
+
+	TupleDesc tupleDescriptor = NULL;
+	Datum values[TABLE_METADATA_FIELDS];
+	bool isNulls[TABLE_METADATA_FIELDS];
+
+	TypeFuncClass resultTypeClass = get_call_result_type(fcinfo, NULL,
+														 &tupleDescriptor);
+	if (resultTypeClass != TYPEFUNC_COMPOSITE)
+	{
+		ereport(ERROR, (errmsg("return type must be a row type")));
+	}
+
+	/* form heap tuple for local disk space statistics */
+	memset(values, 0, sizeof(values));
+	memset(isNulls, false, sizeof(isNulls));
+
+	values[0] = UInt64GetDatum(availableBytes);
+	values[1] = UInt64GetDatum(totalBytes);
+
+	HeapTuple diskSpaceTuple = heap_form_tuple(tupleDescriptor, values, isNulls);
+
+	PG_RETURN_DATUM(HeapTupleGetDatum(diskSpaceTuple));
+}
+
+
+/*
+ * GetLocalDiskSpaceStats returns total and available disk space for the disk containing
+ * PGDATA (not considering tablespaces, quota).
+ */
+static bool
+GetLocalDiskSpaceStats(uint64 *availableBytes, uint64 *totalBytes)
+{
+	struct statvfs buffer;
+	if (statvfs(DataDir, &buffer) != 0)
+	{
+		return false;
+	}
+
+	*availableBytes = buffer.f_bfree * buffer.f_frsize;
+	*totalBytes = buffer.f_blocks * buffer.f_frsize;
+
+	return true;
+}
+
+
 /*
  * citus_database_size is a UDF that returns the sum of the database sizes
  * across all nodes.
@@ -235,7 +393,8 @@ citus_relation_size(PG_FUNCTION_ARGS)
 
 
 /*
- * NodeDatabaseSize returns the size of a database on a given node.
+ * NodeDatabaseSize returns the size of a database on a given node or
+ * 0 in case of failure.
  */
 static uint64
 NodeDatabaseSize(char *nodeName, int nodePort, char *databaseName)
@@ -243,24 +402,37 @@ NodeDatabaseSize(char *nodeName, int nodePort, char *databaseName)
 	uint32 connectionFlags = 0;
 	PGresult *result = NULL;
 	StringInfo sizeQuery = makeStringInfo();
-	bool raiseErrors = true;
+	bool raiseErrors = false;
 
 	appendStringInfo(sizeQuery, "SELECT pg_catalog.pg_database_size(%s)",
 					 quote_literal_cstr(quote_identifier(databaseName)));
 
 	MultiConnection *connection = GetNodeConnection(connectionFlags, nodeName, nodePort);
 
-	int queryResult = ExecuteOptionalRemoteCommand(connection, sizeQuery->data,
+	int queryResult = ExecuteOptionalRemoteCommand(connection, sizeQuery->data, &result);
-												   &result);
 	if (queryResult != 0)
 	{
-		ereport(ERROR, (errcode(ERRCODE_CONNECTION_FAILURE),
+		ereport(WARNING, (errcode(ERRCODE_CONNECTION_FAILURE),
-						errmsg("cannot get the size because of a connection error")));
+						  errmsg("cannot get the database size for node %s:%d",
+								 nodeName, nodePort)));
+		return 0;
 	}
 
-	List *sizeList = ReadFirstColumnAsText(result);
+	if (!IsResponseOK(result) || PQntuples(result) != 1)
-	StringInfo databaseSizeStringInfo = (StringInfo) linitial(sizeList);
+	{
-	uint64 databaseSize = SafeStringToUint64(databaseSizeStringInfo->data);
+		ereport(WARNING, (errcode(ERRCODE_CONNECTION_FAILURE),
+						  errmsg("cannot get the disk space statistics for node %s:%d",
+								 nodeName, nodePort)));
+
+		PQclear(result);
+		ClearResults(connection, raiseErrors);
+
+		return false;
+	}
+
+	char *databaseSizeString = PQgetvalue(result, 0, 0);
+
+	uint64 databaseSize = SafeStringToUint64(databaseSizeString);
 
 	PQclear(result);
 	ClearResults(connection, raiseErrors);

src/backend/distributed/sql/citus--9.5-1--10.0-1.sql

@@ -4,6 +4,8 @@ DROP FUNCTION IF EXISTS pg_catalog.citus_total_relation_size(regclass);
 
 #include "udfs/citus_total_relation_size/10.0-1.sql"
 #include "udfs/citus_tables/10.0-1.sql"
+#include "udfs/citus_local_disk_space_stats/10.0-1.sql"
+#include "udfs/citus_node_disk_space_stats/10.0-1.sql"
 #include "udfs/citus_database_size/10.0-1.sql"
 #include "udfs/citus_node_database_size/10.0-1.sql"
 #include "udfs/citus_nodes/10.0-1.sql"

src/backend/distributed/sql/downgrades/citus--10.0-1--9.5-1.sql

@@ -12,5 +12,7 @@ DROP FUNCTION pg_catalog.citus_total_relation_size(regclass,boolean);
 
 DROP VIEW pg_catalog.citus_nodes;
 
+DROP FUNCTION pg_catalog.citus_local_disk_space_stats();
+DROP FUNCTION pg_catalog.citus_node_disk_space_stats(text,int);
 DROP FUNCTION pg_catalog.citus_database_size(name);
 DROP FUNCTION pg_catalog.citus_node_database_size(text,int,name);

src/backend/distributed/sql/udfs/citus_local_disk_space_stats/10.0-1.sql

@@ -0,0 +1,8 @@
+CREATE OR REPLACE FUNCTION pg_catalog.citus_local_disk_space_stats(
+  OUT available_disk_size bigint,
+  OUT total_disk_size bigint)
+RETURNS record
+LANGUAGE C STRICT
+AS 'MODULE_PATHNAME', $$citus_local_disk_space_stats$$;
+COMMENT ON FUNCTION pg_catalog.citus_local_disk_space_stats()
+IS 'returns statistics on available disk space';

src/backend/distributed/sql/udfs/citus_local_disk_space_stats/latest.sql

@@ -0,0 +1,8 @@
+CREATE OR REPLACE FUNCTION pg_catalog.citus_local_disk_space_stats(
+  OUT available_disk_size bigint,
+  OUT total_disk_size bigint)
+RETURNS record
+LANGUAGE C STRICT
+AS 'MODULE_PATHNAME', $$citus_local_disk_space_stats$$;
+COMMENT ON FUNCTION pg_catalog.citus_local_disk_space_stats()
+IS 'returns statistics on available disk space';

src/backend/distributed/sql/udfs/citus_node_disk_space_stats/10.0-1.sql

@@ -0,0 +1,10 @@
+CREATE OR REPLACE FUNCTION pg_catalog.citus_node_disk_space_stats(
+  nodename text,
+  nodeport int,
+  OUT available_disk_size bigint,
+  OUT total_disk_size bigint)
+RETURNS record
+LANGUAGE C STRICT
+AS 'MODULE_PATHNAME', $$citus_node_disk_space_stats$$;
+COMMENT ON FUNCTION pg_catalog.citus_node_disk_space_stats(text,int)
+IS 'returns statistics on available disk space for the given node';

src/backend/distributed/sql/udfs/citus_node_disk_space_stats/latest.sql

@@ -0,0 +1,10 @@
+CREATE OR REPLACE FUNCTION pg_catalog.citus_node_disk_space_stats(
+  nodename text,
+  nodeport int,
+  OUT available_disk_size bigint,
+  OUT total_disk_size bigint)
+RETURNS record
+LANGUAGE C STRICT
+AS 'MODULE_PATHNAME', $$citus_node_disk_space_stats$$;
+COMMENT ON FUNCTION pg_catalog.citus_node_disk_space_stats(text,int)
+IS 'returns statistics on available disk space for the given node';

src/backend/distributed/sql/udfs/citus_nodes/10.0-1.sql

@@ -31,9 +31,11 @@ SELECT
    USING
      (shardid)
    WHERE p.groupid = n.groupid AND t.partmethod = 'n' AND t.repmodel = 't'
-  ) AS "Reference Tables"
+  ) AS "Reference Tables",
+  round(100 * (1. - (ds.available_disk_size::double precision / ds.total_disk_size))) || '%' AS "Disk Usage"
 FROM
-  pg_dist_node n
+  pg_dist_node n,
+  citus_node_disk_space_stats(n.nodename, n.nodeport) ds
 ORDER BY
   groupid;
 

src/backend/distributed/sql/udfs/citus_nodes/latest.sql

@@ -31,9 +31,11 @@ SELECT
    USING
      (shardid)
    WHERE p.groupid = n.groupid AND t.partmethod = 'n' AND t.repmodel = 't'
-  ) AS "Reference Tables"
+  ) AS "Reference Tables",
+  round(100 * (1. - (ds.available_disk_size::double precision / ds.total_disk_size))) || '%' AS "Disk Usage"
 FROM
-  pg_dist_node n
+  pg_dist_node n,
+  citus_node_disk_space_stats(n.nodename, n.nodeport) ds
 ORDER BY
   groupid;
 

src/test/regress/expected/multi_extension.out

@@ -484,7 +484,9 @@ SELECT * FROM print_extension_changes();
                                               | function alter_columnar_table_set(regclass,integer,integer,name,integer)
                                               | function citus_database_size(name)
                                               | function citus_internal.columnar_ensure_objects_exist()
+                                              | function citus_local_disk_space_stats()
                                               | function citus_node_database_size(text,integer,name)
+                                              | function citus_node_disk_space_stats(text,integer)
                                               | function citus_total_relation_size(regclass,boolean)
                                               | function columnar.columnar_handler(internal)
                                               | schema columnar
@@ -494,7 +496,7 @@ SELECT * FROM print_extension_changes();
                                               | table columnar.options
                                               | view citus_nodes
                                               | view citus_tables
-(16 rows)
+(18 rows)
 
 DROP TABLE prev_objects, extension_diff;
 -- show running version

src/test/regress/expected/multi_extension_0.out

@@ -481,7 +481,9 @@ SELECT * FROM print_extension_changes();
  function citus_total_relation_size(regclass) |
                                               | function citus_database_size(name)
                                               | function citus_internal.columnar_ensure_objects_exist()
+                                              | function citus_local_disk_space_stats()
                                               | function citus_node_database_size(text,integer,name)
+                                              | function citus_node_disk_space_stats(text,integer)
                                               | function citus_total_relation_size(regclass,boolean)
                                               | schema columnar
                                               | sequence columnar.storageid_seq
@@ -490,7 +492,7 @@ SELECT * FROM print_extension_changes();
                                               | table columnar.options
                                               | view citus_nodes
                                               | view citus_tables
-(12 rows)
+(14 rows)
 
 DROP TABLE prev_objects, extension_diff;
 -- show running version

src/test/regress/expected/upgrade_list_citus_objects.out

@@ -51,8 +51,10 @@ ORDER BY 1;
  function citus_json_concatenate_final(json)
  function citus_jsonb_concatenate(jsonb,jsonb)
  function citus_jsonb_concatenate_final(jsonb)
+ function citus_local_disk_space_stats()
  function citus_node_capacity_1(integer)
  function citus_node_database_size(text,integer,name)
+ function citus_node_disk_space_stats(text,integer)
  function citus_prepare_pg_upgrade()
  function citus_query_stats()
  function citus_relation_size(regclass)
@@ -221,5 +223,5 @@ ORDER BY 1;
  view citus_tables
  view citus_worker_stat_activity
  view pg_dist_shard_placement
-(205 rows)
+(207 rows)
 

src/test/regress/expected/upgrade_list_citus_objects_0.out

@@ -48,8 +48,10 @@ ORDER BY 1;
  function citus_json_concatenate_final(json)
  function citus_jsonb_concatenate(jsonb,jsonb)
  function citus_jsonb_concatenate_final(jsonb)
+ function citus_local_disk_space_stats()
  function citus_node_capacity_1(integer)
  function citus_node_database_size(text,integer,name)
+ function citus_node_disk_space_stats(text,integer)
  function citus_prepare_pg_upgrade()
  function citus_query_stats()
  function citus_relation_size(regclass)