Use a lookup table for units in pg_size_pretty and pg_size_bytes

david-rowley · david-rowley · commit 56ff8b29919f · 2021-07-09T16:29:02.000+12:00
We've grown 2 versions of pg_size_pretty over the years, one for BIGINT and one for NUMERIC. Both should output the same, but keeping them in sync is harder than needed due to neither function sharing a source of truth about which units to use and how to transition to the next largest unit. Here we add a static array which defines the units that we recognize and have both pg_size_pretty and pg_size_pretty_numeric use it. This will make adding any units in the future a very simple task. The table contains all information required to allow us to also modify pg_size_bytes to use the lookup table, so adjust that too. There are no behavioral changes here. Author: David Rowley Reviewed-by: Dean Rasheed, Tom Lane, David Christensen Discussion: https://postgr.es/m/CAApHDvru1F7qsEVL-iOHeezJ+5WVxXnyD_Jo9nht+Eh85ekK-Q@mail.gmail.com
diff --git a/src/backend/utils/adt/dbsize.c b/src/backend/utils/adt/dbsize.c
@@ -34,6 +34,27 @@
 /* Divide by two and round away from zero */
 #define half_rounded(x)   (((x) + ((x) < 0 ? -1 : 1)) / 2)
 
+/* Units used in pg_size_pretty functions.  All units must be powers of 2 */
+struct size_pretty_unit
+{
+	const char *name;			/* bytes, kB, MB, GB etc */
+	uint32		limit;			/* upper limit, prior to half rounding after
+								 * converting to this unit. */
+	bool		round;			/* do half rounding for this unit */
+	uint8		unitbits;		/* (1 << unitbits) bytes to make 1 of this
+								 * unit */
+};
+
+/* When adding units here also update the error message in pg_size_bytes */
+static const struct size_pretty_unit size_pretty_units[] = {
+	{"bytes", 10 * 1024, false, 0},
+	{"kB", 20 * 1024 - 1, true, 10},
+	{"MB", 20 * 1024 - 1, true, 20},
+	{"GB", 20 * 1024 - 1, true, 30},
+	{"TB", 20 * 1024 - 1, true, 40},
+	{NULL, 0, false, 0}
+};
+
 /* Return physical size of directory contents, or 0 if dir doesn't exist */
 static int64
 db_dir_size(const char *path)
@@ -535,41 +556,34 @@ pg_size_pretty(PG_FUNCTION_ARGS)
 {
 	int64		size = PG_GETARG_INT64(0);
 	char		buf[64];
-	int64		limit = 10 * 1024;
-	int64		limit2 = limit * 2 - 1;
+	const struct size_pretty_unit *unit;
 
-	if (Abs(size) < limit)
-		snprintf(buf, sizeof(buf), INT64_FORMAT " bytes", size);
-	else
+	for (unit = size_pretty_units; unit->name != NULL; unit++)
 	{
-		/*
-		 * We use divide instead of bit shifting so that behavior matches for
-		 * both positive and negative size values.
-		 */
-		size /= (1 << 9);		/* keep one extra bit for rounding */
-		if (Abs(size) < limit2)
-			snprintf(buf, sizeof(buf), INT64_FORMAT " kB",
-					 half_rounded(size));
-		else
+		uint8		bits;
+
+		/* use this unit if there are no more units or we're below the limit */
+		if (unit[1].name == NULL || Abs(size) < unit->limit)
 		{
-			size /= (1 << 10);
-			if (Abs(size) < limit2)
-				snprintf(buf, sizeof(buf), INT64_FORMAT " MB",
-						 half_rounded(size));
-			else
-			{
-				size /= (1 << 10);
-				if (Abs(size) < limit2)
-					snprintf(buf, sizeof(buf), INT64_FORMAT " GB",
-							 half_rounded(size));
-				else
-				{
-					size /= (1 << 10);
-					snprintf(buf, sizeof(buf), INT64_FORMAT " TB",
-							 half_rounded(size));
-				}
-			}
+			if (unit->round)
+				size = half_rounded(size);
+
+			snprintf(buf, sizeof(buf), INT64_FORMAT " %s", size, unit->name);
+			break;
 		}
+
+		/*
+		 * Determine the number of bits to use to build the divisor.  We may
+		 * need to use 1 bit less than the difference between this and the
+		 * next unit if the next unit uses half rounding.  Or we may need to
+		 * shift an extra bit if this unit uses half rounding and the next one
+		 * does not.  We use division rather than shifting right by this
+		 * number of bits to ensure positive and negative values are rounded
+		 * in the same way.
+		 */
+		bits = (unit[1].unitbits - unit->unitbits - (unit[1].round == true)
+				+ (unit->round == true));
+		size /= ((int64) 1) << bits;
 	}
 
 	PG_RETURN_TEXT_P(cstring_to_text(buf));
@@ -640,57 +654,35 @@ Datum
 pg_size_pretty_numeric(PG_FUNCTION_ARGS)
 {
 	Numeric		size = PG_GETARG_NUMERIC(0);
-	Numeric		limit,
-				limit2;
-	char	   *result;
-
-	limit = int64_to_numeric(10 * 1024);
-	limit2 = int64_to_numeric(10 * 1024 * 2 - 1);
+	char	   *result = NULL;
+	const struct size_pretty_unit *unit;
 
-	if (numeric_is_less(numeric_absolute(size), limit))
+	for (unit = size_pretty_units; unit->name != NULL; unit++)
 	{
-		result = psprintf("%s bytes", numeric_to_cstring(size));
-	}
-	else
-	{
-		/* keep one extra bit for rounding */
-		/* size /= (1 << 9) */
-		size = numeric_truncated_divide(size, 1 << 9);
+		unsigned int shiftby;
 
-		if (numeric_is_less(numeric_absolute(size), limit2))
+		/* use this unit if there are no more units or we're below the limit */
+		if (unit[1].name == NULL ||
+			numeric_is_less(numeric_absolute(size),
+							int64_to_numeric(unit->limit)))
 		{
-			size = numeric_half_rounded(size);
-			result = psprintf("%s kB", numeric_to_cstring(size));
-		}
-		else
-		{
-			/* size /= (1 << 10) */
-			size = numeric_truncated_divide(size, 1 << 10);
-
-			if (numeric_is_less(numeric_absolute(size), limit2))
-			{
+			if (unit->round)
 				size = numeric_half_rounded(size);
-				result = psprintf("%s MB", numeric_to_cstring(size));
-			}
-			else
-			{
-				/* size /= (1 << 10) */
-				size = numeric_truncated_divide(size, 1 << 10);
-
-				if (numeric_is_less(numeric_absolute(size), limit2))
-				{
-					size = numeric_half_rounded(size);
-					result = psprintf("%s GB", numeric_to_cstring(size));
-				}
-				else
-				{
-					/* size /= (1 << 10) */
-					size = numeric_truncated_divide(size, 1 << 10);
-					size = numeric_half_rounded(size);
-					result = psprintf("%s TB", numeric_to_cstring(size));
-				}
-			}
+
+			result = psprintf("%s %s", numeric_to_cstring(size), unit->name);
+			break;
 		}
+
+		/*
+		 * Determine the number of bits to use to build the divisor.  We may
+		 * need to use 1 bit less than the difference between this and the
+		 * next unit if the next unit uses half rounding.  Or we may need to
+		 * shift an extra bit if this unit uses half rounding and the next one
+		 * does not.
+		 */
+		shiftby = (unit[1].unitbits - unit->unitbits - (unit[1].round == true)
+				   + (unit->round == true));
+		size = numeric_truncated_divide(size, ((int64) 1) << shiftby);
 	}
 
 	PG_RETURN_TEXT_P(cstring_to_text(result));
@@ -791,6 +783,7 @@ pg_size_bytes(PG_FUNCTION_ARGS)
 	/* Handle possible unit */
 	if (*strptr != '\0')
 	{
+		const struct size_pretty_unit *unit;
 		int64		multiplier = 0;
 
 		/* Trim any trailing whitespace */
@@ -802,21 +795,18 @@ pg_size_bytes(PG_FUNCTION_ARGS)
 		endptr++;
 		*endptr = '\0';
 
-		/* Parse the unit case-insensitively */
-		if (pg_strcasecmp(strptr, "bytes") == 0)
-			multiplier = (int64) 1;
-		else if (pg_strcasecmp(strptr, "kb") == 0)
-			multiplier = (int64) 1024;
-		else if (pg_strcasecmp(strptr, "mb") == 0)
-			multiplier = ((int64) 1024) * 1024;
-
-		else if (pg_strcasecmp(strptr, "gb") == 0)
-			multiplier = ((int64) 1024) * 1024 * 1024;
-
-		else if (pg_strcasecmp(strptr, "tb") == 0)
-			multiplier = ((int64) 1024) * 1024 * 1024 * 1024;
+		for (unit = size_pretty_units; unit->name != NULL; unit++)
+		{
+			/* Parse the unit case-insensitively */
+			if (pg_strcasecmp(strptr, unit->name) == 0)
+			{
+				multiplier = ((int64) 1) << unit->unitbits;
+				break;
+			}
+		}
 
-		else
+		/* Verify we found a valid unit in the loop above */
+		if (unit->name == NULL)
 			ereport(ERROR,
 					(errcode(ERRCODE_INVALID_PARAMETER_VALUE),
 					 errmsg("invalid size: \"%s\"", text_to_cstring(arg)),
