diff --git a/Makefile b/Makefile index b4e41f3..eac36e1 100644 --- a/Makefile +++ b/Makefile @@ -1,5 +1,5 @@ # It doesn't use CFLAGS in compiling, it uses some postgres magic instead :( -CFLAGS := ${CFLAGS} -Wall -Werror +override CFLAGS = -Wall -Wmissing-prototypes -Wpointer-arith -Wendif-labels -Wmissing-format-attribute -Wformat-security -fno-strict-aliasing -fwrapv -fexcess-precision=standard -Wno-unused-but-set-variable -Wno-address -Wno-format-truncation -Wno-stringop-truncation -g -ggdb -std=gnu99 -Werror=uninitialized -Werror=implicit-function-declaration -DGPBUILD EXTENSION = query_tag EXTVERSION = $(shell grep default_version $(EXTENSION).control | \ diff --git a/sql/query_tag--1.0.sql b/sql/query_tag--1.0.sql index f5e874b..c7a544a 100644 --- a/sql/query_tag--1.0.sql +++ b/sql/query_tag--1.0.sql @@ -6,7 +6,7 @@ CREATE FUNCTION CURRENT_RESGROUP() RETURNS text AS 'query_tag', 'current_resgroup' LANGUAGE C STRICT VOLATILE; -CREATE OR REPLACE FUNCTION is_tag_in_guc(query_tag text) +CREATE FUNCTION is_tag_in_guc(text) RETURNS boolean AS 'query_tag', 'is_tag_in_guc' LANGUAGE C STRICT VOLATILE; @@ -32,7 +32,8 @@ CREATE TABLE wlm_rules ( spill_file_mb int, cpuskew_percent int, cpuskew_duration_sec int, - order_id int not null + order_id int not null, + kill_rule boolean ); CREATE TABLE gpcc_wlm_log_history ( diff --git a/src/parser.c b/src/parser.c index 02e7b4d..28ff8b5 100644 --- a/src/parser.c +++ b/src/parser.c @@ -4,188 +4,84 @@ #include "utils/builtins.h" #include "c.h" #include +#include "utils/memutils.h" -static bool good_char(char c) { - return (isalnum(c) || c == '=' || c == ';' || c == '_'); -} - -bool parse_tags(const char *guc, List **parsed) { - *parsed = NIL; - bool is_in_key = true; - bool is_empty = true; - const char *substr_ptr_start; - const char *substr_ptr_end; - char *key_tag = NULL, *value_tag = NULL; - if (!guc || *guc == '\0') { - return true; +static void free_double_list(List **double_list) { + ListCell *current_cell; + foreach(current_cell, *double_list) { + list_free(lfirst(current_cell)); } - substr_ptr_start = guc; - const char *cur; - for (cur = guc;; ++cur) { - if (*cur == '\0') { - if (is_in_key) { - return false; - } - if (is_empty) { - return true; - } - substr_ptr_end = cur; - value_tag = copy_substr(substr_ptr_start, substr_ptr_end); - substr_ptr_start = cur + 1; - - tag_pair *tpair; - tpair = malloc(sizeof(tpair)); - tpair->key = key_tag; - tpair->value = value_tag; - *parsed = lappend(*parsed, tpair); - break; - } - if (!good_char(*cur)) { - return false; - } - if (*cur == '=') { - if (is_empty) { - return false; - } - if (is_in_key) { - is_in_key = false; - is_empty = true; - } else { - return false; - } - } - if (*cur == ';') { - if (is_empty) { - return false; - } - if (is_in_key) { - return false; - } else { - is_in_key = true; - is_empty = true; - } - } - if (isalpha(*cur) || *cur == '_') { - is_empty = false; - } - if (*cur == '=') { - substr_ptr_end = cur; - key_tag = copy_substr(substr_ptr_start, substr_ptr_end); - substr_ptr_start = cur + 1; - } - if (*cur == ';') { - substr_ptr_end = cur; - value_tag = copy_substr(substr_ptr_start, substr_ptr_end); - substr_ptr_start = cur + 1; - - tag_pair *tpair; - tpair = malloc(sizeof(tpair)); - tpair->key = key_tag; - tpair->value = value_tag; - *parsed = lappend(*parsed, tpair); - } - } - - return true; + list_free(*double_list); + *double_list = NIL; } -void tag_list_free(List **tag_list) { - ListCell *cell; - foreach (cell, *tag_list) { - tag_pair *tpair = lfirst(cell); - free(tpair->key); - free(tpair->value); - free(tpair); +bool split_tags(const char *tags, ParsedTags **parsed) { + if (!tags) { + elog(ERROR, "QUERY_TAG: there's a NULL query_tag somewhere."); } - list_free(*tag_list); - *tag_list = NIL; -} - -bool is_safe(const char *guc) { - if (!guc) + MemoryContext oldctx = MemoryContextSwitchTo(TopMemoryContext); + *parsed = palloc(sizeof(**parsed)); + char *tags_mutable = pstrdup(tags); + List *tag_pairs = NIL; + bool ok = SplitIdentifierString(tags_mutable, ';', &tag_pairs); + if (!ok) { + pfree(*parsed); + *parsed = NULL; + pfree(tags_mutable); + list_free(tag_pairs); + MemoryContextSwitchTo(oldctx); return false; - if (*guc == '\0') - return true; - bool is_in_key = true; - bool is_empty = true; - for (const char *cur = guc; *cur != '\0'; ++cur) { - if (!good_char(*cur)) { - elog(INFO, "GUHed char: %c", *cur); + } + ListCell *current_tag; + List *parsed_tags = NIL; + foreach(current_tag, tag_pairs) { + List *parsed_current_tag = NIL; + ok = SplitIdentifierString(lfirst(current_tag), '=', &parsed_current_tag); + if (!ok || list_length(parsed_current_tag) != 2) { + list_free(parsed_current_tag); + free_double_list(&parsed_tags); + pfree(*parsed); + *parsed = NULL; + pfree(tags_mutable); + list_free(tag_pairs); + MemoryContextSwitchTo(oldctx); return false; } - if (*cur == '=') { - if (is_empty) { - return false; - } - if (is_in_key) { - is_in_key = false; - is_empty = true; - } else { - return false; - } - } - if (*cur == ';') { - if (is_empty) { - return false; - } - if (is_in_key) { - return false; - } else { - is_in_key = true; - is_empty = true; - } - } - if (isalnum(*cur) || *cur == '_') { - is_empty = false; - } + parsed_tags = lappend(parsed_tags, parsed_current_tag); } - if (is_in_key && !is_empty) - return false; - return true; -} -char *copy_substr(const char *start, const char *end) { - int len = end - start; - char *ret = calloc(len + 1, sizeof(char)); - memcpy(ret, start, len); - ret[len] = '\0'; - return ret; + (*parsed)->tags_mutable = tags_mutable; + (*parsed)->parsed_tags = parsed_tags; + MemoryContextSwitchTo(oldctx); + return true; } -bool is_tag_in_guc_ctype(const char *tag, const char *guc) { - List *parsed_guc = NIL; - List *parsed_tag = NIL; - - bool ok = parse_tags(guc, &parsed_guc); - if (!ok) { - tag_list_free(&parsed_guc); - return false; - } - - ok = parse_tags(tag, &parsed_tag); - if (!ok) { - tag_list_free(&parsed_tag); - return false; - } - ListCell *guc_pair_cell, *tag_pair_cell; - foreach (tag_pair_cell, parsed_tag) { - tag_pair *cur_tag = lfirst(tag_pair_cell); +bool is_parsed_rule_in_parsed_guc(ParsedTags *rule_tags, ParsedTags *guc_tags) { + ListCell *guc_cell, *tag_cell; + foreach (tag_cell, rule_tags->parsed_tags) { + List *tag_pair = lfirst(tag_cell); bool found = false; - foreach (guc_pair_cell, parsed_guc) { - tag_pair *cur_guc = lfirst(guc_pair_cell); - if (strcmp(cur_tag->key, cur_guc->key) == 0 && - strcmp(cur_tag->value, cur_guc->value) == 0) { + foreach (guc_cell, guc_tags->parsed_tags) { + List *guc_pair = lfirst(guc_cell); + if (strcmp(linitial(guc_pair), linitial(tag_pair)) == 0 && + strcmp(lsecond(guc_pair), lsecond(tag_pair)) == 0) { found = true; break; } } if (!found) { - tag_list_free(&parsed_guc); - tag_list_free(&parsed_tag); return false; } } - tag_list_free(&parsed_guc); - tag_list_free(&parsed_tag); return true; } + +void free_parsed_tags(ParsedTags **parsed) { + if (*parsed) { + free_double_list(&(*parsed)->parsed_tags); + if ((*parsed)->tags_mutable) + pfree((*parsed)->tags_mutable); + pfree(*parsed); + *parsed = NULL; + } +} \ No newline at end of file diff --git a/src/parser.h b/src/parser.h index 5cc26ec..571fb5d 100644 --- a/src/parser.h +++ b/src/parser.h @@ -1,12 +1,10 @@ #include "nodes/pg_list.h" -typedef struct tpair{ - char *key; - char *value; -} tag_pair; +typedef struct ParsedTags { + char *tags_mutable; + List *parsed_tags; +} ParsedTags; -bool is_safe(const char *guc); -bool parse_tags(const char *guc, List **parsed); -void tag_list_free(List **tag_list); -char *copy_substr(const char *start, const char *end); -bool is_tag_in_guc_ctype(const char *tag, const char *guc); +bool split_tags(const char *tags, ParsedTags **parsed); +bool is_parsed_rule_in_parsed_guc(ParsedTags *parsed_rule, ParsedTags *parsed_guc); +void free_parsed_tags(ParsedTags **parsed); diff --git a/src/query_tag.c b/src/query_tag.c index 901ccef..14add84 100644 --- a/src/query_tag.c +++ b/src/query_tag.c @@ -1,5 +1,7 @@ #include +#include +#include #include #include #include @@ -12,6 +14,9 @@ #include "parser.h" #include "utils/syscache.h" #include "catalog/pg_authid.h" +#include "cdb/cdbvars.h" +#include "optimizer/planner.h" +#include PG_MODULE_MAGIC; @@ -25,28 +30,81 @@ static bool check_new_query_tag(char **, void **, GucSource); void _PG_init(void); void _PG_fini(void); +static ParsedTags *parsed_guc_tags; static char *query_tag = NULL; const static char *NO_GROUP_MSG = "unknown"; static const int MAX_QUERY_SIZE = 200; static const int MAX_QUERY_TAG_LENGTH = 100; static resgroup_assign_hook_type prev_hook = NULL; +static planner_hook_type prev_planner_hook = NULL; + +PlannedStmt * +kill_rules_manager(Query *parse, int cursorOptions, ParamListInfo boundParams) +{ + PlannedStmt *result = NULL; + result = prev_planner_hook(parse, cursorOptions, boundParams); + return result; + + // under construction + + int n_moves = result->nMotionNodes; + int SPI_status = -1; + StringInfoData query; + StringInfoData cost_query; + char *rgname = GetResGroupNameForId(current_resgroup_id()); + char *rolename = GetUserNameFromId(GetUserId()); + initStringInfo(&query); + initStringInfo(&cost_query); + double cost = result->planTree->total_cost; + elog(NOTICE, "Cost: %f", cost); + if (result->planGen == PLANGEN_PLANNER) { + appendStringInfo(&cost_query, "planner_cost <= %f", cost); + } else { + appendStringInfo(&cost_query, "orca_cost <= %f", cost); + } + appendStringInfo(&query, "select rule_id, dest_resg from wlm_rules where resgname = '%s' and role " + "= '%s' and active = TRUE and %s and is_tag_in_guc(query_tag) and kill_rule = TRUE order by order_id limit 1;", + rgname, rolename, cost_query.data); + PG_TRY(); + { + SPI_status = SPI_execute(query.data, false, 0); + } + PG_CATCH(); + { + SPI_status = -1; + } + PG_END_TRY(); -Datum is_tag_in_guc(PG_FUNCTION_ARGS) { - text *query_tag = PG_GETARG_TEXT_P(0); - char *query_tag_str = text_to_cstring(query_tag); - - const char *guc_query_tag_value = - GetConfigOption("QUERY_TAG", false, false); - bool result = is_tag_in_guc_ctype(query_tag_str, guc_query_tag_value); - - pfree(query_tag_str); + if (SPI_status < 0) { + SPI_finish(); + elog(DEBUG3, "QUERY_TAG: query failed in the kill rule processor."); + return result; + } + if (SPI_processed > 0) { + elog(ERROR, "QUERY_TAG: there's a kill_rule. Stopped"); + } + return result; +} +Datum is_tag_in_guc(PG_FUNCTION_ARGS) { + text *rule_query_tag = PG_GETARG_TEXT_P(0); + char *rule_query_tag_cstr = text_to_cstring(rule_query_tag); + if (!rule_query_tag_cstr) { + PG_RETURN_BOOL(false); + } + ParsedTags *parsed_rule_tags = NULL; + bool ok = split_tags(rule_query_tag_cstr, &parsed_rule_tags); + if (!ok) { + pfree(rule_query_tag_cstr); + PG_RETURN_BOOL(false); + } + bool result = is_parsed_rule_in_parsed_guc(parsed_rule_tags, parsed_guc_tags); + free_parsed_tags(&parsed_rule_tags); PG_RETURN_BOOL(result); } static Oid resgroup_assign_by_query_tag(void) { - Oid groupId = InvalidOid; if (prev_hook) groupId = prev_hook(); @@ -94,7 +152,7 @@ static Oid resgroup_assign_by_query_tag(void) { int full_length = snprintf( query, sizeof(query), "select rule_id, dest_resg from wlm_rules where resgname = '%s' and role " - "= '%s' and active = TRUE and is_tag_in_guc(query_tag) order by order_id limit 1;", + "= '%s' and active = TRUE and is_tag_in_guc(query_tag) and kill_rule = FALSE order by order_id limit 1;", rgname, rolename); if (full_length >= MAX_QUERY_SIZE) { elog(ERROR, "QUERY_TAG: failed, query tag too long"); @@ -114,7 +172,7 @@ static Oid resgroup_assign_by_query_tag(void) { SPI_finish(); return groupId; } - elog(DEBUG3, "QUERY_TAG: set resgroup to: %s, %d", crgname, + elog(DEBUG1, "QUERY_TAG: set resgroup to: %s, %d", crgname, GetResGroupIdForName(crgname)); SPI_finish(); return GetResGroupIdForName(crgname); @@ -125,15 +183,32 @@ static Oid resgroup_assign_by_query_tag(void) { } static bool check_new_query_tag(char **newvalue, void **extra, GucSource source) { + if (!*newvalue) { + free_parsed_tags(&parsed_guc_tags); + parsed_guc_tags = NULL; + return true; + } + if (query_tag && (strcmp(*newvalue, query_tag) == 0)) { + return true; + } elog(DEBUG3, "QUERY_TAG: Checking new tag"); if (strlen(*newvalue) >= MAX_QUERY_TAG_LENGTH) { elog(DEBUG3, "QUERY_TAG: Tag too long, didn't set."); return false; } - if (!is_safe(*newvalue)) { + ParsedTags *new_parsed_guc_tag = NULL; + bool ok = split_tags(*newvalue, &new_parsed_guc_tag); + if (!ok) { + elog(DEBUG3, "QUERY_TAG: Tag didn't pass parsing."); return false; } - elog(DEBUG3, "QUERY_TAG: All okay, set new tag."); + // if they are the same, it means, that memory contexts messed up + assert(parsed_guc_tags != new_parsed_guc_tag); + if (parsed_guc_tags) { + free_parsed_tags(&parsed_guc_tags); + parsed_guc_tags = NULL; + } + parsed_guc_tags = new_parsed_guc_tag; return true; } @@ -160,11 +235,17 @@ void _PG_init(void) { &query_tag, "", /* initial value */ PGC_USERSET, 0, /* flags */ - check_new_query_tag, /* check hook */ + check_new_query_tag, /* check hook */ NULL, /* assign hook */ NULL); /* show hook */ prev_hook = resgroup_assign_hook; + if (planner_hook) { + prev_planner_hook = planner_hook; + } else { + prev_planner_hook = standard_planner; + } resgroup_assign_hook = resgroup_assign_by_query_tag; + planner_hook = kill_rules_manager; } void _PG_fini(void) { diff --git a/test/expected/query_tag.out b/test/expected/query_tag.out index 775aebe..3fe691f 100644 --- a/test/expected/query_tag.out +++ b/test/expected/query_tag.out @@ -18,18 +18,19 @@ SHOW QUERY_TAG; (1 row) SET QUERY_TAG TO "a=a;"; +ERROR: invalid value for parameter "QUERY_TAG": "a=a;" SHOW QUERY_TAG; - QUERY_TAG ------------ - a=a; + QUERY_TAG +------------- + aboba=aboba (1 row) SET QUERY_TAG TO "a;"; ERROR: invalid value for parameter "QUERY_TAG": "a;" SHOW QUERY_TAG; - QUERY_TAG ------------ - a=a; + QUERY_TAG +------------- + aboba=aboba (1 row) SET QUERY_TAG TO "123=123"; @@ -50,13 +51,11 @@ SHOW QUERY_TAG; -- complex behaviour -- should set to rgroup1 CREATE RESOURCE GROUP rgroup1 WITH (CPU_RATE_LIMIT=20, MEMORY_LIMIT=25, MEMORY_SPILL_RATIO=20); -ERROR: resource group "rgroup1" already exists CREATE RESOURCE GROUP rgroup2 WITH (CPU_RATE_LIMIT=20, MEMORY_LIMIT=25, MEMORY_SPILL_RATIO=20); -ERROR: resource group "rgroup2" already exists -INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active) - VALUES ('admin_group', 'clepip', 'rgroup1', 3, 'group=rgroup1', TRUE); -INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active) - VALUES ('admin_group', 'clepip', 'rgroup2', 4, 'group=rgroup1', TRUE); +INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active, kill_rule) + VALUES ('admin_group', 'clepip', 'rgroup1', 3, 'group=rgroup1', TRUE, FALSE); +INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active, kill_rule) + VALUES ('admin_group', 'clepip', 'rgroup2', 4, 'group=rgroup1', TRUE, FALSE); SET QUERY_TAG TO "group=rgroup1"; SELECT current_resgroup(); current_resgroup @@ -66,10 +65,10 @@ SELECT current_resgroup(); -- should still set to rgroup1 (we are in admin_group, the "current_rsgroup" shows group after --- applying tags, and tags are applied accordingly to users rsgroup) -INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active) - VALUES ('rgroup1', 'clepip', 'rgroup2', 1, 'group=still_1', TRUE); -INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active) - VALUES ('admin_group', 'clepip', 'rgroup1', 5, 'group=still_1', TRUE); +INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active, kill_rule) + VALUES ('rgroup1', 'clepip', 'rgroup2', 1, 'group=still_1', TRUE, FALSE); +INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active, kill_rule) + VALUES ('admin_group', 'clepip', 'rgroup1', 5, 'group=still_1', TRUE, FALSE); SET QUERY_TAG TO "group=still_1"; SELECT CURRENT_RESGROUP(); current_resgroup @@ -78,10 +77,10 @@ SELECT CURRENT_RESGROUP(); (1 row) -- should set to rgroup1 -INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active) - VALUES ('admin_group', 'clepip', 'rgroup2', 1, 'group=cancelled', FALSE); -INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active) - VALUES ('admin_group', 'clepip', 'rgroup1', 2, 'group=cancelled', TRUE); +INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active, kill_rule) + VALUES ('admin_group', 'clepip', 'rgroup2', 1, 'group=cancelled', FALSE, FALSE); +INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active, kill_rule) + VALUES ('admin_group', 'clepip', 'rgroup1', 2, 'group=cancelled', TRUE, FALSE); SET QUERY_TAG TO "group=cancelled"; SELECT CURRENT_RESGROUP(); current_resgroup diff --git a/test/sql/query_tag.sql b/test/sql/query_tag.sql index d157553..134a1a7 100644 --- a/test/sql/query_tag.sql +++ b/test/sql/query_tag.sql @@ -30,10 +30,10 @@ SHOW QUERY_TAG; CREATE RESOURCE GROUP rgroup1 WITH (CPU_RATE_LIMIT=20, MEMORY_LIMIT=25, MEMORY_SPILL_RATIO=20); CREATE RESOURCE GROUP rgroup2 WITH (CPU_RATE_LIMIT=20, MEMORY_LIMIT=25, MEMORY_SPILL_RATIO=20); -INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active) - VALUES ('admin_group', 'clepip', 'rgroup1', 3, 'group=rgroup1', TRUE); -INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active) - VALUES ('admin_group', 'clepip', 'rgroup2', 4, 'group=rgroup1', TRUE); +INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active, kill_rule) + VALUES ('admin_group', 'clepip', 'rgroup1', 3, 'group=rgroup1', TRUE, FALSE); +INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active, kill_rule) + VALUES ('admin_group', 'clepip', 'rgroup2', 4, 'group=rgroup1', TRUE, FALSE); SET QUERY_TAG TO "group=rgroup1"; SELECT current_resgroup(); @@ -41,20 +41,20 @@ SELECT current_resgroup(); -- should still set to rgroup1 (we are in admin_group, the "current_rsgroup" shows group after --- applying tags, and tags are applied accordingly to users rsgroup) -INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active) - VALUES ('rgroup1', 'clepip', 'rgroup2', 1, 'group=still_1', TRUE); -INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active) - VALUES ('admin_group', 'clepip', 'rgroup1', 5, 'group=still_1', TRUE); +INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active, kill_rule) + VALUES ('rgroup1', 'clepip', 'rgroup2', 1, 'group=still_1', TRUE, FALSE); +INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active, kill_rule) + VALUES ('admin_group', 'clepip', 'rgroup1', 5, 'group=still_1', TRUE, FALSE); SET QUERY_TAG TO "group=still_1"; SELECT CURRENT_RESGROUP(); -- should set to rgroup1 -INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active) - VALUES ('admin_group', 'clepip', 'rgroup2', 1, 'group=cancelled', FALSE); -INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active) - VALUES ('admin_group', 'clepip', 'rgroup1', 2, 'group=cancelled', TRUE); +INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active, kill_rule) + VALUES ('admin_group', 'clepip', 'rgroup2', 1, 'group=cancelled', FALSE, FALSE); +INSERT INTO wlm_rules (resgname, role, dest_resg, order_id, query_tag, active, kill_rule) + VALUES ('admin_group', 'clepip', 'rgroup1', 2, 'group=cancelled', TRUE, FALSE); SET QUERY_TAG TO "group=cancelled"; SELECT CURRENT_RESGROUP(); \ No newline at end of file