From 8b387db81b93349968be3ae9d6cd7cd77d05d2e7 Mon Sep 17 00:00:00 2001 From: David Rowley Date: Mon, 5 Jul 2021 20:03:27 +1200 Subject: [PATCH v1] Use lookup table for units in pg_size_pretty and pg_size_bytes 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 would be harder than needed due to neither function sharing a source of truth about which units to use. Here we add a static array that 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 also use it there too, so adjust that too. There are no behavioral changes here. --- src/backend/utils/adt/dbsize.c | 153 +++++++++++++++------------------ 1 file changed, 68 insertions(+), 85 deletions(-) diff --git a/src/backend/utils/adt/dbsize.c b/src/backend/utils/adt/dbsize.c index 9c39e7d3b3..4b7331d85c 100644 --- a/src/backend/utils/adt/dbsize.c +++ b/src/backend/utils/adt/dbsize.c @@ -34,6 +34,29 @@ /* Divide by two and round towards positive infinity. */ #define half_rounded(x) (((x) + ((x) < 0 ? 0 : 1)) / 2) +/* Units used in pg_size_pretty functions. All units must be powers of 2 */ +struct size_pretty_unit +{ + char *text; /* bytes, kB, MB, GB etc */ + uint32 limit; /* upper limit after reducing by shiftRight + * bits from previous unit */ + uint8 shiftRight; /* number of bits to shift right by after + * limit is exceeded before considering the + * next unit. Remaining shifts will be half + * rounded until we get to unitBits */ + uint8 unitBits; /* 1 << unitBits 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, 9, 0}, /* keep one extra bit for rounding */ + {"kB", 20 * 1024 - 1, 10, 10}, + {"MB", 20 * 1024 - 1, 10, 20}, + {"GB", 20 * 1024 - 1, 10, 30}, + {"TB", 20 * 1024 - 1, 10, 40}, + {NULL, 0, 0, 0} +}; + /* Return physical size of directory contents, or 0 if dir doesn't exist */ static int64 db_dir_size(const char *path) @@ -535,37 +558,25 @@ pg_size_pretty(PG_FUNCTION_ARGS) { int64 size = PG_GETARG_INT64(0); char buf[64]; - int64 limit = 10 * 1024; - int64 limit2 = limit * 2 - 1; + uint8 rightshifts = 0; + 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->text !=NULL; unit++) { - size >>= 9; /* keep one extra bit for rounding */ - if (Abs(size) < limit2) - snprintf(buf, sizeof(buf), INT64_FORMAT " kB", - half_rounded(size)); - else + /* use this unit if below limit or there's no more units */ + if (Abs(size) < unit->limit || unit[1].text == NULL) { - size >>= 10; - if (Abs(size) < limit2) - snprintf(buf, sizeof(buf), INT64_FORMAT " MB", - half_rounded(size)); - else - { - size >>= 10; - if (Abs(size) < limit2) - snprintf(buf, sizeof(buf), INT64_FORMAT " GB", - half_rounded(size)); - else - { - size >>= 10; - snprintf(buf, sizeof(buf), INT64_FORMAT " TB", - half_rounded(size)); - } - } + /* half-round until we get down to unitBits */ + while (rightshifts++ < unit->unitBits) + size = half_rounded(size); + + snprintf(buf, sizeof(buf), INT64_FORMAT " %s", size, unit->text); + break; } + + /* size too big for this unit, shift right and try the next unit */ + size >>= unit->shiftRight; + rightshifts += unit->shiftRight; } PG_RETURN_TEXT_P(cstring_to_text(buf)); @@ -636,56 +647,30 @@ 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); + Numeric limit; + char *result = NULL; + uint8 rightshifts = 0; + const struct size_pretty_unit *unit; - if (numeric_is_less(numeric_absolute(size), limit)) - { - result = psprintf("%s bytes", numeric_to_cstring(size)); - } - else + for (unit = size_pretty_units; unit->text !=NULL; unit++) { - /* keep one extra bit for rounding */ - /* size >>= 9 */ - size = numeric_shift_right(size, 9); + limit = int64_to_numeric(unit->limit); - if (numeric_is_less(numeric_absolute(size), limit2)) - { - size = numeric_half_rounded(size); - result = psprintf("%s kB", numeric_to_cstring(size)); - } - else + /* use this unit if below limit or there's no more units */ + if (numeric_is_less(numeric_absolute(size), limit) || + unit[1].text == NULL) { - /* size >>= 10 */ - size = numeric_shift_right(size, 10); - if (numeric_is_less(numeric_absolute(size), limit2)) - { + /* half-round until we get down to unitBits */ + while (rightshifts++ < unit->unitBits) size = numeric_half_rounded(size); - result = psprintf("%s MB", numeric_to_cstring(size)); - } - else - { - /* size >>= 10 */ - size = numeric_shift_right(size, 10); - - if (numeric_is_less(numeric_absolute(size), limit2)) - { - size = numeric_half_rounded(size); - result = psprintf("%s GB", numeric_to_cstring(size)); - } - else - { - /* size >>= 10 */ - size = numeric_shift_right(size, 10); - size = numeric_half_rounded(size); - result = psprintf("%s TB", numeric_to_cstring(size)); - } - } + + result = psprintf("%s %s", numeric_to_cstring(size), unit->text); + break; } + + /* size too big for this unit, shift right and try the next unit */ + size = numeric_shift_right(size, unit->shiftRight); + rightshifts += unit->shiftRight; } PG_RETURN_TEXT_P(cstring_to_text(result)); @@ -786,6 +771,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 */ @@ -797,21 +783,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->text !=NULL; unit++) + { + /* Parse the unit case-insensitively */ + if (pg_strcasecmp(strptr, unit->text) == 0) + { + multiplier = ((int64) 1) << unit->unitBits; + break; + } + } - else + /* Verify we found a valid unit in the loop above */ + if (unit->text == NULL) ereport(ERROR, (errcode(ERRCODE_INVALID_PARAMETER_VALUE), errmsg("invalid size: \"%s\"", text_to_cstring(arg)), -- 2.30.2