diff --git a/fdw/Makefile b/fdw/Makefile index 12d07308a..1d8099c5d 100644 --- a/fdw/Makefile +++ b/fdw/Makefile @@ -1,7 +1,7 @@ EXTENSION = pxf_fdw -DATA = pxf_fdw--2.0.sql pxf_fdw--1.0--2.0.sql pxf_fdw--2.0--1.0.sql pxf_fdw--1.0.sql +DATA = pxf_fdw--2.1.sql pxf_fdw--2.0--2.1.sql pxf_fdw--2.1--2.0.sql pxf_fdw--2.0.sql pxf_fdw--1.0--2.0.sql pxf_fdw--2.0--1.0.sql pxf_fdw--1.0.sql MODULE_big = pxf_fdw -OBJS = pxf_fdw.o pxf_bridge.o pxf_deparse.o pxf_filter.o pxf_header.o pxf_option.o libchurl.o +OBJS = pxf_fdw.o pxf_bridge.o pxf_deparse.o pxf_filter.o pxf_header.o pxf_option.o pxf_stat_activity.o pxf_cancel_activity.o libchurl.o REGRESS = pxf_fdw_wrapper pxf_fdw_server pxf_fdw_user_mapping pxf_fdw_foreign_table SHLIB_LINK += -lcurl diff --git a/fdw/pxf_cancel_activity.c b/fdw/pxf_cancel_activity.c new file mode 100644 index 000000000..eb28551c8 --- /dev/null +++ b/fdw/pxf_cancel_activity.c @@ -0,0 +1,190 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/* + * pxf_cancel_activity.c + * + * Backs the pxf_cancel_backend and pxf_interrupt_backend SQL functions. Each is + * dispatched with EXECUTE ON ALL SEGMENTS: every segment asks its local PXF + * instance to terminate the in-flight requests of a given Greenplum session + * that originate from its own segment id (passed via X-GP-SEGMENT-ID). Because a + * single PXF instance serves every segment co-located on the host, this + * per-segment filter guarantees each running request is acted on exactly once, + * by its owning segment. + * + * Both functions are set-returning (one row per segment, carrying the PXF + * response verbatim as a JSON text value) because EXECUTE ON ALL SEGMENTS is + * only permitted for set-returning functions; the counts are summed by the SQL + * wrappers, keeping this layer schema-agnostic. + */ + +#include "libchurl.h" + +#include "fmgr.h" +#include "funcapi.h" +#include "lib/stringinfo.h" +#include "cdb/cdbvars.h" +#include "utils/builtins.h" + +/* PXF service defaults, mirror of pxf_stat_activity.c */ +#define PXF_CANCEL_DEFAULT_HOST "localhost" +#define PXF_CANCEL_DEFAULT_PORT "5888" +#define PXF_CANCEL_SERVICE_PREFIX "pxf" +#define PXF_CANCEL_READ_BUFFER_SIZE (64 * 1024) + +PG_FUNCTION_INFO_V1(pxf_cancel_backend_raw); +PG_FUNCTION_INFO_V1(pxf_interrupt_backend_raw); + +Datum pxf_cancel_backend_raw(PG_FUNCTION_ARGS); +Datum pxf_interrupt_backend_raw(PG_FUNCTION_ARGS); + +/* + * Returns the PXF service authority (host:port) for the local instance, + * honoring the PXF_HOST / PXF_PORT environment variables and falling back to + * localhost:5888, consistent with how segments reach their local PXF. + */ +static char * +get_pxf_cancel_authority(void) +{ + char *host = getenv("PXF_HOST"); + char *port = getenv("PXF_PORT"); + + return psprintf("%s:%s", + host ? host : PXF_CANCEL_DEFAULT_HOST, + port ? port : PXF_CANCEL_DEFAULT_PORT); +} + +/* + * Issues an HTTP GET to the given local PXF endpoint, scoped to this segment and + * the target Greenplum session, and returns the raw JSON response body as a + * palloc'd text value in the current memory context. + */ +static text * +fetch_backend_control_body(const char *endpoint, int32 session_id) +{ + CHURL_HEADERS headers; + CHURL_HANDLE handle; + StringInfoData uri; + StringInfoData response; + char readbuf[PXF_CANCEL_READ_BUFFER_SIZE]; + char segment_id[32]; + char session_buf[32]; + size_t n; + text *result; + + /* build the request URI for the local PXF instance */ + initStringInfo(&uri); + appendStringInfo(&uri, "http://%s/%s/%s", + get_pxf_cancel_authority(), PXF_CANCEL_SERVICE_PREFIX, endpoint); + + /* scope the request to this segment so co-located segments don't act twice */ + snprintf(segment_id, sizeof(segment_id), "%d", GpIdentity.segindex); + snprintf(session_buf, sizeof(session_buf), "%d", session_id); + + headers = churl_headers_init(); + churl_headers_append(headers, "X-GP-SEGMENT-ID", segment_id); + churl_headers_append(headers, "X-GP-SESSION-ID", session_buf); + churl_headers_append(headers, "Accept", "application/json"); + + elog(DEBUG2, "%s: segment %d requesting %s for session %d", + endpoint, GpIdentity.segindex, uri.data, session_id); + + handle = churl_init_download(uri.data, headers); + + /* read the full JSON body */ + initStringInfo(&response); + while ((n = churl_read(handle, readbuf, sizeof(readbuf))) != 0) + appendBinaryStringInfo(&response, readbuf, n); + + /* surface any error reported by the PXF service on the closed connection */ + churl_read_check_connectivity(handle); + + churl_cleanup(handle, false); + churl_headers_cleanup(headers); + + result = cstring_to_text(response.data); + + pfree(uri.data); + pfree(response.data); + + return result; +} + +/* + * Shared set-returning body: emit a single row carrying the raw JSON response + * from the local PXF instance for the given endpoint and session id. Dispatched + * with EXECUTE ON ALL SEGMENTS (which Cloudberry only permits for set-returning + * functions); the pxf_cancel_backend / pxf_interrupt_backend SQL wrappers sum + * the per-segment counts. + */ +static Datum +backend_control_srf(FunctionCallInfo fcinfo, const char *endpoint) +{ + FuncCallContext *funcctx; + + if (SRF_IS_FIRSTCALL()) + { + MemoryContext oldcontext; + int32 session_id = PG_GETARG_INT32(0); + + funcctx = SRF_FIRSTCALL_INIT(); + + /* the fetched body must outlive this call, so build it in the + * multi-call context and hand it back one row at a time */ + oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx); + + funcctx->user_fctx = fetch_backend_control_body(endpoint, session_id); + funcctx->max_calls = 1; + + MemoryContextSwitchTo(oldcontext); + } + + funcctx = SRF_PERCALL_SETUP(); + + if (funcctx->call_cntr < funcctx->max_calls) + SRF_RETURN_NEXT(funcctx, PointerGetDatum((text *) funcctx->user_fctx)); + + SRF_RETURN_DONE(funcctx); +} + +/* + * pxf_cancel_backend_raw(session_id int) + * + * Asks the local PXF instance to gracefully cancel the in-flight requests of the + * given session (by ending their current bridge). Returns the raw JSON body, + * e.g. {"cancelled":N}. + */ +Datum +pxf_cancel_backend_raw(PG_FUNCTION_ARGS) +{ + return backend_control_srf(fcinfo, "cancel_backend"); +} + +/* + * pxf_interrupt_backend_raw(session_id int) + * + * Asks the local PXF instance to interrupt the worker thread(s) of the in-flight + * requests of the given session. Returns the raw JSON body, e.g. + * {"interrupted":N}. + */ +Datum +pxf_interrupt_backend_raw(PG_FUNCTION_ARGS) +{ + return backend_control_srf(fcinfo, "interrupt_backend"); +} diff --git a/fdw/pxf_fdw--2.0--2.1.sql b/fdw/pxf_fdw--2.0--2.1.sql new file mode 100644 index 000000000..aee9c9fa0 --- /dev/null +++ b/fdw/pxf_fdw--2.0--2.1.sql @@ -0,0 +1,87 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +/* fdw/pxf_fdw--2.0--2.1.sql */ + +-- complain if script is sourced in psql, rather than via ALTER EXTENSION +\echo Use "ALTER EXTENSION pxf_fdw UPDATE TO '2.1'" to load this file. \quit + +-- Raw per-segment accessor: each segment asks its local PXF instance for the +-- activity that originates from its own segment id and returns the JSON body +-- verbatim as a single row. Dispatched to every segment; the typed columns are +-- produced by the pxf_stat_activity view below. The function is set-returning +-- (one row per segment) because EXECUTE ON ALL SEGMENTS is only permitted for +-- set-returning functions. +CREATE FUNCTION pxf_stat_activity_raw() RETURNS SETOF text +AS 'MODULE_PATHNAME', 'pxf_stat_activity_raw' +LANGUAGE C VOLATILE EXECUTE ON ALL SEGMENTS; + +-- pg_stat_activity-like view of the queries currently running inside PXF, +-- aggregated across all segment hosts. DISTINCT is a safety net; PXF already +-- de-duplicates by filtering each response to the requesting segment id. +CREATE VIEW pxf_stat_activity AS +SELECT DISTINCT + (a->>'segmentId')::int AS segment_id, + (a->>'gpSessionId')::int AS session_id, + (a->>'gpCommandCount')::int AS command_count, + a->>'transactionId' AS xid, + a->>'requestType' AS operation, + a->>'user' AS usename, + a->>'serverName' AS server, + a->>'profile' AS profile, + a->>'schemaName' AS schema_name, + a->>'tableName' AS table_name, + a->>'dataSource' AS data_source, + to_timestamp((a->>'startTimeMs')::bigint / 1000.0) AS query_start, + a->>'host' AS pxf_host +FROM ( + SELECT json_array_elements(raw::json -> 'activities') AS a + FROM pxf_stat_activity_raw() AS raw +) s; + +-- Per-segment cancellation primitives backing pxf_cancel_backend / +-- pxf_interrupt_backend. Each segment asks its local PXF instance to terminate +-- the in-flight requests of the given Cloudberry session that originate from its +-- own segment id, returning the JSON body verbatim as a single row (e.g. +-- {"cancelled":N} / {"interrupted":N}). Dispatched to every segment; the counts +-- are summed by the SQL wrappers below. Set-returning because EXECUTE ON ALL +-- SEGMENTS is only permitted for set-returning functions. +CREATE FUNCTION pxf_cancel_backend_raw(session_id int) RETURNS SETOF text +AS 'MODULE_PATHNAME', 'pxf_cancel_backend_raw' +LANGUAGE C VOLATILE STRICT EXECUTE ON ALL SEGMENTS; + +CREATE FUNCTION pxf_interrupt_backend_raw(session_id int) RETURNS SETOF text +AS 'MODULE_PATHNAME', 'pxf_interrupt_backend_raw' +LANGUAGE C VOLATILE STRICT EXECUTE ON ALL SEGMENTS; + +-- Gracefully cancels the in-flight PXF requests of a Cloudberry session across +-- the whole cluster by ending their current bridge. Returns the number of +-- requests that were signalled. Analogous to pg_cancel_backend, but keyed by +-- the Cloudberry session id (as reported in pxf_stat_activity.session_id). +CREATE FUNCTION pxf_cancel_backend(session_id int) RETURNS int AS $$ + SELECT coalesce(sum((raw::json ->> 'cancelled')::int), 0)::int + FROM pxf_cancel_backend_raw(session_id) AS raw +$$ LANGUAGE sql VOLATILE; + +-- Interrupts the worker thread(s) of the in-flight PXF requests of a Cloudberry +-- session across the whole cluster. Returns the number of requests that were +-- interrupted. A forceful complement to pxf_cancel_backend for requests that do +-- not observe cancellation (e.g. blocked in a non-interruptible read). +CREATE FUNCTION pxf_interrupt_backend(session_id int) RETURNS int AS $$ + SELECT coalesce(sum((raw::json ->> 'interrupted')::int), 0)::int + FROM pxf_interrupt_backend_raw(session_id) AS raw +$$ LANGUAGE sql VOLATILE; diff --git a/fdw/pxf_fdw--2.1--2.0.sql b/fdw/pxf_fdw--2.1--2.0.sql new file mode 100644 index 000000000..a4e2c2b5b --- /dev/null +++ b/fdw/pxf_fdw--2.1--2.0.sql @@ -0,0 +1,28 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +/* fdw/pxf_fdw--2.1--2.0.sql */ + +-- complain if script is sourced in psql, rather than via ALTER EXTENSION +\echo Use "ALTER EXTENSION pxf_fdw UPDATE TO '2.0'" to load this file. \quit + +DROP FUNCTION pxf_interrupt_backend(int); +DROP FUNCTION pxf_cancel_backend(int); +DROP FUNCTION pxf_interrupt_backend_raw(int); +DROP FUNCTION pxf_cancel_backend_raw(int); +DROP VIEW pxf_stat_activity; +DROP FUNCTION pxf_stat_activity_raw(); diff --git a/fdw/pxf_fdw--2.1.sql b/fdw/pxf_fdw--2.1.sql new file mode 100644 index 000000000..496d39d2d --- /dev/null +++ b/fdw/pxf_fdw--2.1.sql @@ -0,0 +1,142 @@ +-- Licensed to the Apache Software Foundation (ASF) under one +-- or more contributor license agreements. See the NOTICE file +-- distributed with this work for additional information +-- regarding copyright ownership. The ASF licenses this file +-- to you under the Apache License, Version 2.0 (the +-- "License"); you may not use this file except in compliance +-- with the License. You may obtain a copy of the License at +-- +-- http://www.apache.org/licenses/LICENSE-2.0 +-- +-- Unless required by applicable law or agreed to in writing, +-- software distributed under the License is distributed on an +-- "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +-- KIND, either express or implied. See the License for the +-- specific language governing permissions and limitations +-- under the License. + +/* fdw/pxf_fdw--2.1.sql */ + +-- complain if script is sourced in psql, rather than via CREATE EXTENSION +\echo Use "CREATE EXTENSION pxf_fdw" to load this file. \quit + +CREATE FUNCTION pxf_fdw_handler() +RETURNS fdw_handler +AS 'MODULE_PATHNAME' +LANGUAGE C STRICT; + +CREATE FUNCTION pxf_fdw_validator(text[], oid) +RETURNS void +AS 'MODULE_PATHNAME' +LANGUAGE C STRICT; + +CREATE FOREIGN DATA WRAPPER jdbc_pxf_fdw + HANDLER pxf_fdw_handler + VALIDATOR pxf_fdw_validator + OPTIONS ( protocol 'jdbc', mpp_execute 'all segments' ); + +CREATE FOREIGN DATA WRAPPER hdfs_pxf_fdw + HANDLER pxf_fdw_handler + VALIDATOR pxf_fdw_validator + OPTIONS ( protocol 'hdfs', mpp_execute 'all segments' ); + +CREATE FOREIGN DATA WRAPPER hive_pxf_fdw + HANDLER pxf_fdw_handler + VALIDATOR pxf_fdw_validator + OPTIONS ( protocol 'hive', mpp_execute 'all segments' ); + +CREATE FOREIGN DATA WRAPPER hbase_pxf_fdw + HANDLER pxf_fdw_handler + VALIDATOR pxf_fdw_validator + OPTIONS ( protocol 'hbase', mpp_execute 'all segments' ); + +CREATE FOREIGN DATA WRAPPER s3_pxf_fdw + HANDLER pxf_fdw_handler + VALIDATOR pxf_fdw_validator + OPTIONS ( protocol 's3', mpp_execute 'all segments' ); + +CREATE FOREIGN DATA WRAPPER gs_pxf_fdw + HANDLER pxf_fdw_handler + VALIDATOR pxf_fdw_validator + OPTIONS ( protocol 'gs', mpp_execute 'all segments' ); + +CREATE FOREIGN DATA WRAPPER abfss_pxf_fdw + HANDLER pxf_fdw_handler + VALIDATOR pxf_fdw_validator + OPTIONS ( protocol 'abfss', mpp_execute 'all segments' ); + +CREATE FOREIGN DATA WRAPPER wasbs_pxf_fdw + HANDLER pxf_fdw_handler + VALIDATOR pxf_fdw_validator + OPTIONS ( protocol 'wasbs', mpp_execute 'all segments' ); + +CREATE FOREIGN DATA WRAPPER file_pxf_fdw + HANDLER pxf_fdw_handler + VALIDATOR pxf_fdw_validator + OPTIONS ( protocol 'file', mpp_execute 'all segments' ); + +-- Raw per-segment accessor: each segment asks its local PXF instance for the +-- activity that originates from its own segment id and returns the JSON body +-- verbatim as a single row. Dispatched to every segment; the typed columns are +-- produced by the pxf_stat_activity view below. The function is set-returning +-- (one row per segment) because EXECUTE ON ALL SEGMENTS is only permitted for +-- set-returning functions. +CREATE FUNCTION pxf_stat_activity_raw() RETURNS SETOF text +AS 'MODULE_PATHNAME', 'pxf_stat_activity_raw' +LANGUAGE C VOLATILE EXECUTE ON ALL SEGMENTS; + +-- pg_stat_activity-like view of the queries currently running inside PXF, +-- aggregated across all segment hosts. DISTINCT is a safety net; PXF already +-- de-duplicates by filtering each response to the requesting segment id. +CREATE VIEW pxf_stat_activity AS +SELECT DISTINCT + (a->>'segmentId')::int AS segment_id, + (a->>'gpSessionId')::int AS session_id, + (a->>'gpCommandCount')::int AS command_count, + a->>'transactionId' AS xid, + a->>'requestType' AS operation, + a->>'user' AS usename, + a->>'serverName' AS server, + a->>'profile' AS profile, + a->>'schemaName' AS schema_name, + a->>'tableName' AS table_name, + a->>'dataSource' AS data_source, + to_timestamp((a->>'startTimeMs')::bigint / 1000.0) AS query_start, + a->>'host' AS pxf_host +FROM ( + SELECT json_array_elements(raw::json -> 'activities') AS a + FROM pxf_stat_activity_raw() AS raw +) s; + +-- Per-segment cancellation primitives backing pxf_cancel_backend / +-- pxf_interrupt_backend. Each segment asks its local PXF instance to terminate +-- the in-flight requests of the given Cloudberry session that originate from its +-- own segment id, returning the JSON body verbatim as a single row (e.g. +-- {"cancelled":N} / {"interrupted":N}). Dispatched to every segment; the counts +-- are summed by the SQL wrappers below. Set-returning because EXECUTE ON ALL +-- SEGMENTS is only permitted for set-returning functions. +CREATE FUNCTION pxf_cancel_backend_raw(session_id int) RETURNS SETOF text +AS 'MODULE_PATHNAME', 'pxf_cancel_backend_raw' +LANGUAGE C VOLATILE STRICT EXECUTE ON ALL SEGMENTS; + +CREATE FUNCTION pxf_interrupt_backend_raw(session_id int) RETURNS SETOF text +AS 'MODULE_PATHNAME', 'pxf_interrupt_backend_raw' +LANGUAGE C VOLATILE STRICT EXECUTE ON ALL SEGMENTS; + +-- Gracefully cancels the in-flight PXF requests of a Cloudberry session across +-- the whole cluster by ending their current bridge. Returns the number of +-- requests that were signalled. Analogous to pg_cancel_backend, but keyed by +-- the Cloudberry session id (as reported in pxf_stat_activity.session_id). +CREATE FUNCTION pxf_cancel_backend(session_id int) RETURNS int AS $$ + SELECT coalesce(sum((raw::json ->> 'cancelled')::int), 0)::int + FROM pxf_cancel_backend_raw(session_id) AS raw +$$ LANGUAGE sql VOLATILE; + +-- Interrupts the worker thread(s) of the in-flight PXF requests of a Cloudberry +-- session across the whole cluster. Returns the number of requests that were +-- interrupted. A forceful complement to pxf_cancel_backend for requests that do +-- not observe cancellation (e.g. blocked in a non-interruptible read). +CREATE FUNCTION pxf_interrupt_backend(session_id int) RETURNS int AS $$ + SELECT coalesce(sum((raw::json ->> 'interrupted')::int), 0)::int + FROM pxf_interrupt_backend_raw(session_id) AS raw +$$ LANGUAGE sql VOLATILE; diff --git a/fdw/pxf_fdw.control b/fdw/pxf_fdw.control index e10226c73..5026d7db9 100644 --- a/fdw/pxf_fdw.control +++ b/fdw/pxf_fdw.control @@ -1,6 +1,6 @@ # pxf_fdw extension directory = 'extension' comment = 'PXF Foreign Data Wrapper for Greenplum' -default_version = '2.0' +default_version = '2.1' module_pathname = '$libdir/pxf_fdw' relocatable = true diff --git a/fdw/pxf_stat_activity.c b/fdw/pxf_stat_activity.c new file mode 100644 index 000000000..3b21b0774 --- /dev/null +++ b/fdw/pxf_stat_activity.c @@ -0,0 +1,159 @@ +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +/* + * pxf_stat_activity.c + * + * Backs the pxf_stat_activity SQL view. The function pxf_stat_activity_raw() + * is dispatched with EXECUTE ON ALL SEGMENTS: each segment asks its local PXF + * instance for the requests that originate from its own segment id (passed via + * the X-GP-SEGMENT-ID header). Because a single PXF instance serves every + * segment co-located on the host, this per-segment filter guarantees each + * running query is reported exactly once, by its owning segment. + * + * The function is set-returning (one row per segment, carrying the PXF response + * verbatim as a JSON text value) because EXECUTE ON ALL SEGMENTS is only + * permitted for set-returning functions; the typed columns are produced by the + * SQL view, keeping this layer schema-agnostic. + */ + +#include "libchurl.h" + +#include "fmgr.h" +#include "funcapi.h" +#include "lib/stringinfo.h" +#include "cdb/cdbvars.h" +#include "utils/builtins.h" + +/* PXF service defaults, mirror of pxf_option.h / external-table pxfutils.h */ +#define PXF_STAT_DEFAULT_HOST "localhost" +#define PXF_STAT_DEFAULT_PORT "5888" +#define PXF_STAT_SERVICE_PREFIX "pxf" +#define PXF_STAT_READ_BUFFER_SIZE (64 * 1024) + +PG_FUNCTION_INFO_V1(pxf_stat_activity_raw); + +Datum pxf_stat_activity_raw(PG_FUNCTION_ARGS); + +/* + * Returns the PXF service authority (host:port) for the local instance, + * honoring the PXF_HOST / PXF_PORT environment variables and falling back to + * localhost:5888, consistent with how segments reach their local PXF. + */ +static char * +get_pxf_stat_authority(void) +{ + char *host = getenv("PXF_HOST"); + char *port = getenv("PXF_PORT"); + + return psprintf("%s:%s", + host ? host : PXF_STAT_DEFAULT_HOST, + port ? port : PXF_STAT_DEFAULT_PORT); +} + +/* + * Issues an HTTP GET to the local PXF /pxf/stat_activity endpoint, scoped to + * this segment, and returns the raw JSON response body as a palloc'd text value + * in the current memory context. + */ +static text * +fetch_stat_activity_body(void) +{ + CHURL_HEADERS headers; + CHURL_HANDLE handle; + StringInfoData uri; + StringInfoData response; + char readbuf[PXF_STAT_READ_BUFFER_SIZE]; + char segment_id[32]; + size_t n; + text *result; + + /* build the request URI for the local PXF instance */ + initStringInfo(&uri); + appendStringInfo(&uri, "http://%s/%s/stat_activity", + get_pxf_stat_authority(), PXF_STAT_SERVICE_PREFIX); + + /* scope the request to this segment so co-located segments don't duplicate */ + snprintf(segment_id, sizeof(segment_id), "%d", GpIdentity.segindex); + + headers = churl_headers_init(); + churl_headers_append(headers, "X-GP-SEGMENT-ID", segment_id); + churl_headers_append(headers, "Accept", "application/json"); + + elog(DEBUG2, "pxf_stat_activity: segment %d requesting %s", + GpIdentity.segindex, uri.data); + + handle = churl_init_download(uri.data, headers); + + /* read the full JSON body */ + initStringInfo(&response); + while ((n = churl_read(handle, readbuf, sizeof(readbuf))) != 0) + appendBinaryStringInfo(&response, readbuf, n); + + /* surface any error reported by the PXF service on the closed connection */ + churl_read_check_connectivity(handle); + + churl_cleanup(handle, false); + churl_headers_cleanup(headers); + + result = cstring_to_text(response.data); + + pfree(uri.data); + pfree(response.data); + + return result; +} + +/* + * pxf_stat_activity_raw + * + * Set-returning function dispatched with EXECUTE ON ALL SEGMENTS (which + * Cloudberry only permits for set-returning functions). Each segment emits a + * single row: the raw JSON response body from its local PXF instance, scoped to + * that segment. The pxf_stat_activity view unions these rows and expands them + * into typed columns. + */ +Datum +pxf_stat_activity_raw(PG_FUNCTION_ARGS) +{ + FuncCallContext *funcctx; + + if (SRF_IS_FIRSTCALL()) + { + MemoryContext oldcontext; + + funcctx = SRF_FIRSTCALL_INIT(); + + /* the fetched body must outlive this call, so build it in the + * multi-call context and hand it back one row at a time */ + oldcontext = MemoryContextSwitchTo(funcctx->multi_call_memory_ctx); + + funcctx->user_fctx = fetch_stat_activity_body(); + funcctx->max_calls = 1; + + MemoryContextSwitchTo(oldcontext); + } + + funcctx = SRF_PERCALL_SETUP(); + + if (funcctx->call_cntr < funcctx->max_calls) + SRF_RETURN_NEXT(funcctx, PointerGetDatum((text *) funcctx->user_fctx)); + + SRF_RETURN_DONE(funcctx); +} diff --git a/server/pxf-service/src/main/java/org/apache/cloudberry/pxf/service/activity/ActiveRequest.java b/server/pxf-service/src/main/java/org/apache/cloudberry/pxf/service/activity/ActiveRequest.java new file mode 100644 index 000000000..37f9c5919 --- /dev/null +++ b/server/pxf-service/src/main/java/org/apache/cloudberry/pxf/service/activity/ActiveRequest.java @@ -0,0 +1,167 @@ +package org.apache.cloudberry.pxf.service.activity; + +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import lombok.extern.slf4j.Slf4j; +import org.greenplum.pxf.service.bridge.Bridge; + +import java.util.concurrent.locks.ReentrantLock; + +/** + * A single in-flight PXF request tracked by {@link ActiveRequestRegistry}. + *

+ * In addition to the immutable {@link ActiveRequestInfo} snapshot (which defines + * the `pxf_stat_activity` JSON wire contract) it holds the two live + * handles used to terminate the request from another thread: + * - the worker `thread` that is processing it (for `pxf_interrupt_backend`) + * - the current `bridge` (for `pxf_cancel_backend`) + */ +@Slf4j +class ActiveRequest { + + /* + * Concurrency: PXF processes a request synchronously on a servlet + * thread that is returned to a pool afterwards, so a naive + * `thread.interrupt()` could land on an unrelated request the thread has + * since picked up. To prevent that, every mutation that races with request + * completion is serialized on `lock` and gated by the `finished` + * flag: + *

+ */ + + /** Immutable snapshot exposed via {@code pxf_stat_activity}. */ + final ActiveRequestInfo info; + + /** The servlet worker thread that is processing this request. */ + private final Thread thread; + + /** Serializes cancellation/interruption against request completion. */ + private final ReentrantLock lock = new ReentrantLock(); + + /** Set once the request has completed; guarded by {@link #lock}. */ + private boolean finished; + + /** Set when the request has been asked to stop; read by the worker loop. */ + private volatile boolean cancelled; + + /** The bridge currently in use by the worker; guarded by {@link #lock}. */ + private Bridge bridge; + + ActiveRequest(ActiveRequestInfo info, Thread thread) { + this.info = info; + this.thread = thread; + } + + /** + * @return whether this request has been asked to stop; polled by the worker + * loop so it can break out between fragments/records without waiting for a + * blocking read to be unblocked by {@link #cancelIfActive()}. + */ + boolean isCancelled() { + return cancelled; + } + + /** + * Records (or clears, with {@code null}) the bridge the worker is currently + * iterating over, so that {@link #cancelIfActive()} can end it. + */ + void setBridge(Bridge bridge) { + lock.lock(); + try { + this.bridge = bridge; + } finally { + lock.unlock(); + } + } + + /** + * Marks the request as completed. After this returns, {@link #cancelIfActive()} + * and {@link #interruptIfActive()} become no-ops, so the worker thread can be + * safely returned to the pool without risking a stray interrupt. + */ + void markFinished() { + lock.lock(); + try { + finished = true; + } finally { + lock.unlock(); + } + } + + /** + * Requests cancellation of this request if it is still active: raises the + * {@link #cancelled} flag and ends the current bridge to unblock a pending + * read. Safe to call concurrently with the worker's own bridge lifecycle; + * {@code endIteration()} must therefore be idempotent enough to tolerate a + * double close (its error is logged and swallowed). + * + * @return {@code true} if the request was still active and was signalled + */ + boolean cancelIfActive() { + lock.lock(); + try { + if (finished) { + return false; + } + cancelled = true; + if (bridge != null) { + try { + bridge.endIteration(); + } catch (Exception e) { + log.warn("Ignoring error while cancelling bridge for {}", describe(), e); + } + } + return true; + } finally { + lock.unlock(); + } + } + + /** + * Interrupts the worker thread if the request is still active. Holding the + * lock while checking {@link #finished} guarantees the thread has not yet + * returned to the pool, so the interrupt cannot leak onto a later request. + * + * @return {@code true} if the request was still active and was interrupted + */ + boolean interruptIfActive() { + lock.lock(); + try { + if (finished) { + return false; + } + thread.interrupt(); + return true; + } finally { + lock.unlock(); + } + } + + private String describe() { + return String.format("session %d, segment %d, xid %s", + info.getGpSessionId(), info.getSegmentId(), info.getTransactionId()); + } +} diff --git a/server/pxf-service/src/main/java/org/apache/cloudberry/pxf/service/activity/ActiveRequestInfo.java b/server/pxf-service/src/main/java/org/apache/cloudberry/pxf/service/activity/ActiveRequestInfo.java new file mode 100644 index 000000000..ecde7b161 --- /dev/null +++ b/server/pxf-service/src/main/java/org/apache/cloudberry/pxf/service/activity/ActiveRequestInfo.java @@ -0,0 +1,96 @@ +package org.apache.cloudberry.pxf.service.activity; + +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import lombok.Getter; +import org.greenplum.pxf.api.model.RequestContext; + +/** + * Immutable snapshot of an in-flight PXF request, captured at the moment the + * request started being processed. Instances are exposed via the + * `/pxf/stat_activity` endpoint. + *

+ * The getters are serialized to JSON by Jackson. + */ +@Getter +public class ActiveRequestInfo { + + /** The Cloudberry segment id that originated the request. */ + private final int segmentId; + + /** The Cloudberry session id (aka ssid). */ + private final int gpSessionId; + + /** The Cloudberry command count within the session (aka ccnt). */ + private final int gpCommandCount; + + /** The Cloudberry transaction id (XID) of the originating query. */ + private final String transactionId; + + /** The kind of operation: READ_BRIDGE or WRITE_BRIDGE. */ + private final String requestType; + + /** The end-user identity that issued the request. */ + private final String user; + + /** The name of the PXF server configuration used by the request. */ + private final String serverName; + + /** The profile associated with the request (e.g. hdfs:text). */ + private final String profile; + + /** The originating Cloudberry schema name. */ + private final String schemaName; + + /** The originating Cloudberry table name. */ + private final String tableName; + + /** The data source (file path or external resource identifier). */ + private final String dataSource; + + /** Epoch milliseconds when the request started being processed. */ + private final long startTimeMs; + + /** The hostname of the PXF instance that is serving the request. */ + private final String host; + + /** + * Captures a snapshot of the relevant fields from the request context. + * + * @param context the request context of the in-flight request + * @param startTimeMs epoch milliseconds when processing started + * @param host hostname of the PXF instance serving the request + */ + public ActiveRequestInfo(RequestContext context, long startTimeMs, String host) { + this.segmentId = context.getSegmentId(); + this.gpSessionId = context.getGpSessionId(); + this.gpCommandCount = context.getGpCommandCount(); + this.transactionId = context.getTransactionId(); + this.requestType = context.getRequestType() == null ? null : context.getRequestType().name(); + this.user = context.getUser(); + this.serverName = context.getServerName(); + this.profile = context.getProfile(); + this.schemaName = context.getSchemaName(); + this.tableName = context.getTableName(); + this.dataSource = context.getDataSource(); + this.startTimeMs = startTimeMs; + this.host = host; + } +} diff --git a/server/pxf-service/src/main/java/org/apache/cloudberry/pxf/service/activity/ActiveRequestRegistry.java b/server/pxf-service/src/main/java/org/apache/cloudberry/pxf/service/activity/ActiveRequestRegistry.java new file mode 100644 index 000000000..8c1379709 --- /dev/null +++ b/server/pxf-service/src/main/java/org/apache/cloudberry/pxf/service/activity/ActiveRequestRegistry.java @@ -0,0 +1,191 @@ +package org.apache.cloudberry.pxf.service.activity; + +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import lombok.extern.slf4j.Slf4j; +import org.greenplum.pxf.api.model.RequestContext; +import org.greenplum.pxf.service.bridge.Bridge; +import org.springframework.stereotype.Component; + +import java.net.InetAddress; +import java.net.UnknownHostException; +import java.util.ArrayList; +import java.util.List; +import java.util.Set; +import java.util.concurrent.ConcurrentHashMap; + +/** + * Registry of the requests that are currently being processed by this PXF instance. + * A request is registered when it starts being processed by + * `BaseServiceImpl.processData()` and unregistered when it ends. + */ +@Component +@Slf4j +public class ActiveRequestRegistry { + + /** Sentinel meaning "all segments" for the segment-id filter. */ + public static final int ALL_SEGMENTS = -1; + + /** All in-flight requests on this instance. */ + private final Set activeRequests = ConcurrentHashMap.newKeySet(); + + /** + * The request currently being processed by the calling worker thread, used + * by the worker itself `attachBridge`, `isCurrentCancelled`) and to locate + * the entry to remove on `unregister` without threading a handle through the + * processing call chain. Set on `register` and cleared on `unregister`. + */ + private final ThreadLocal currentRequest = new ThreadLocal<>(); + + private final String hostname; + + public ActiveRequestRegistry() { + this.hostname = resolveHostname(); + } + + /** + * Registers the request being processed by the calling thread as in-flight. + * Must be paired with a call to `unregister()` in a finally block, on the + * same thread. + * + * @param context the request context of the request being processed + */ + public void register(RequestContext context) { + ActiveRequestInfo info = new ActiveRequestInfo(context, System.currentTimeMillis(), hostname); + ActiveRequest request = new ActiveRequest(info, Thread.currentThread()); + activeRequests.add(request); + currentRequest.set(request); + } + + /** + * Removes the request registered by the calling thread from the registry and + * marks it finished so that any concurrent `cancel()` / `interrupt` can no + * longer act on the (soon to be recycled) worker thread. No-op if the calling + * thread has no registered request. + */ + public void unregister() { + ActiveRequest request = currentRequest.get(); + if (request != null) { + activeRequests.remove(request); + request.markFinished(); + } + currentRequest.remove(); + } + + /** + * Records the bridge the calling worker thread is currently iterating over, + * so that a concurrent `cancel` can end it. Pass `null` to clear + * the reference once the bridge is closed. No-op if the caller is not a + * registered worker thread. + * + * @param bridge the bridge in use, or `null` to detach + */ + public void attachBridge(Bridge bridge) { + ActiveRequest request = currentRequest.get(); + if (request != null) { + request.setBridge(bridge); + } + } + + /** + * @return whether the request being processed by the calling worker thread + * has been asked to cancel. Polled by the read/write loops so they stop + * between fragments/records. + */ + public boolean isCurrentCancelled() { + ActiveRequest request = currentRequest.get(); + return request != null && request.isCancelled(); + } + + /** + * Returns a snapshot of the currently active requests, optionally filtered + * by the originating segment id. + * + * @param segmentId the originating segment id to filter by, or + * `ALL_SEGMENTS` to return activity for all segments + * @return list of active request descriptors + */ + public List snapshot(int segmentId) { + List result = new ArrayList<>(activeRequests.size()); + for (ActiveRequest request : activeRequests) { + if (matchesSegment(request, segmentId)) { + result.add(request.info); + } + } + return result; + } + + /** + * Gracefully cancels every active request of the given Cloudberry session on + * this segment by ending its current bridge (see `ActiveRequest.cancelIfActive()`). + * + * @param segmentId the originating segment id, or `ALL_SEGMENTS` + * @param sessionId the Cloudberry session id whose requests to cancel + * @return the number of active requests that were signalled + */ + public int cancel(int segmentId, int sessionId) { + int count = 0; + for (ActiveRequest request : activeRequests) { + if (matchesSegment(request, segmentId) && request.info.getGpSessionId() == sessionId) { + if (request.cancelIfActive()) { + count++; + } + } + } + log.info("pxf_cancel_backend: signalled {} request(s) for session {} on segment {}", + count, sessionId, segmentId); + return count; + } + + /** + * Interrupts the worker thread of every active request of the given + * Cloudberry session on this segment (see `ActiveRequest.interruptIfActive()`). + * + * @param segmentId the originating segment id, or `ALL_SEGMENTS` + * @param sessionId the Cloudberry session id whose requests to interrupt + * @return the number of active requests that were interrupted + */ + public int interrupt(int segmentId, int sessionId) { + int count = 0; + for (ActiveRequest request : activeRequests) { + if (matchesSegment(request, segmentId) && request.info.getGpSessionId() == sessionId) { + if (request.interruptIfActive()) { + count++; + } + } + } + log.info("pxf_interrupt_backend: interrupted {} request(s) for session {} on segment {}", + count, sessionId, segmentId); + return count; + } + + private static boolean matchesSegment(ActiveRequest request, int segmentId) { + return segmentId == ALL_SEGMENTS || request.info.getSegmentId() == segmentId; + } + + private static String resolveHostname() { + try { + return InetAddress.getLocalHost().getHostName(); + } catch (UnknownHostException e) { + log.warn("Unable to resolve local hostname for pxf_stat_activity reporting", e); + return "unknown"; + } + } +} diff --git a/server/pxf-service/src/main/java/org/apache/cloudberry/pxf/service/rest/PxfBackendControlResource.java b/server/pxf-service/src/main/java/org/apache/cloudberry/pxf/service/rest/PxfBackendControlResource.java new file mode 100644 index 000000000..c44d0c016 --- /dev/null +++ b/server/pxf-service/src/main/java/org/apache/cloudberry/pxf/service/rest/PxfBackendControlResource.java @@ -0,0 +1,83 @@ +package org.apache.cloudberry.pxf.service.rest; + +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import org.apache.cloudberry.pxf.service.activity.ActiveRequestRegistry; +import org.springframework.http.MediaType; +import org.springframework.http.ResponseEntity; +import org.springframework.web.bind.annotation.GetMapping; +import org.springframework.web.bind.annotation.RequestHeader; +import org.springframework.web.bind.annotation.RequestMapping; +import org.springframework.web.bind.annotation.RestController; + +import java.util.Collections; +import java.util.Map; + +/** + * REST endpoints that terminate in-flight PXF requests, backing the + * `pxf_cancel_backend` and `pxf_interrupt_backend` SQL functions. + */ +@RestController +@RequestMapping("/pxf") +public class PxfBackendControlResource { + + private static final String SEGMENT_ID_HEADER = "X-GP-SEGMENT-ID"; + private static final String SESSION_ID_HEADER = "X-GP-SESSION-ID"; + + private final ActiveRequestRegistry activeRequestRegistry; + + public PxfBackendControlResource(ActiveRequestRegistry activeRequestRegistry) { + this.activeRequestRegistry = activeRequestRegistry; + } + + /** + * Gracefully cancels the active requests of the given session on this segment + * by ending their current bridge. + * + * @param sessionId the Cloudberry session id whose requests to cancel + * @param segmentId the originating segment id; when absent, all segments + * @return a JSON object `{"cancelled":N}` with the number signalled + */ + @GetMapping(value = "/cancel_backend", produces = MediaType.APPLICATION_JSON_VALUE) + public ResponseEntity> cancelBackend( + @RequestHeader(SESSION_ID_HEADER) int sessionId, + @RequestHeader(value = SEGMENT_ID_HEADER, required = false) Integer segmentId) { + int filter = (segmentId == null) ? ActiveRequestRegistry.ALL_SEGMENTS : segmentId; + int cancelled = activeRequestRegistry.cancel(filter, sessionId); + return ResponseEntity.ok(Collections.singletonMap("cancelled", cancelled)); + } + + /** + * Interrupts the worker thread(s) of the active requests of the given session + * on this segment. + * + * @param sessionId the Cloudberry session id whose requests to interrupt + * @param segmentId the originating segment id; when absent, all segments + * @return a JSON object `{"interrupted":N}` with the number interrupted + */ + @GetMapping(value = "/interrupt_backend", produces = MediaType.APPLICATION_JSON_VALUE) + public ResponseEntity> interruptBackend( + @RequestHeader(SESSION_ID_HEADER) int sessionId, + @RequestHeader(value = SEGMENT_ID_HEADER, required = false) Integer segmentId) { + int filter = (segmentId == null) ? ActiveRequestRegistry.ALL_SEGMENTS : segmentId; + int interrupted = activeRequestRegistry.interrupt(filter, sessionId); + return ResponseEntity.ok(Collections.singletonMap("interrupted", interrupted)); + } +} diff --git a/server/pxf-service/src/main/java/org/apache/cloudberry/pxf/service/rest/PxfStatActivityResource.java b/server/pxf-service/src/main/java/org/apache/cloudberry/pxf/service/rest/PxfStatActivityResource.java new file mode 100644 index 000000000..94c1c0320 --- /dev/null +++ b/server/pxf-service/src/main/java/org/apache/cloudberry/pxf/service/rest/PxfStatActivityResource.java @@ -0,0 +1,66 @@ +package org.apache.cloudberry.pxf.service.rest; + +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import org.apache.cloudberry.pxf.service.activity.ActiveRequestInfo; +import org.apache.cloudberry.pxf.service.activity.ActiveRequestRegistry; +import org.springframework.http.MediaType; +import org.springframework.http.ResponseEntity; +import org.springframework.web.bind.annotation.GetMapping; +import org.springframework.web.bind.annotation.RequestHeader; +import org.springframework.web.bind.annotation.RequestMapping; +import org.springframework.web.bind.annotation.RestController; + +import java.util.Collections; +import java.util.List; +import java.util.Map; + +/** + * REST endpoint that exposes the requests currently being processed by this PXF + * instance. It backs the `pxf_stat_activity` SQL view. + */ +@RestController +@RequestMapping("/pxf") +public class PxfStatActivityResource { + + private static final String SEGMENT_ID_HEADER = "X-GP-SEGMENT-ID"; + + private final ActiveRequestRegistry activeRequestRegistry; + + public PxfStatActivityResource(ActiveRequestRegistry activeRequestRegistry) { + this.activeRequestRegistry = activeRequestRegistry; + } + + /** + * Returns the active requests known to this PXF instance, optionally + * filtered by the originating segment id supplied in the request header. + * + * @param segmentId the originating segment id; when absent, activity for all + * segments is returned + * @return a JSON object `{"activities":[...]}` + */ + @GetMapping(value = "/stat_activity", produces = MediaType.APPLICATION_JSON_VALUE) + public ResponseEntity>> statActivity( + @RequestHeader(value = SEGMENT_ID_HEADER, required = false) Integer segmentId) { + int filter = (segmentId == null) ? ActiveRequestRegistry.ALL_SEGMENTS : segmentId; + List activities = activeRequestRegistry.snapshot(filter); + return ResponseEntity.ok(Collections.singletonMap("activities", activities)); + } +} diff --git a/server/pxf-service/src/main/java/org/greenplum/pxf/service/PxfServiceApplication.java b/server/pxf-service/src/main/java/org/greenplum/pxf/service/PxfServiceApplication.java index cacafb4a7..fb4f24ac3 100644 --- a/server/pxf-service/src/main/java/org/greenplum/pxf/service/PxfServiceApplication.java +++ b/server/pxf-service/src/main/java/org/greenplum/pxf/service/PxfServiceApplication.java @@ -12,7 +12,7 @@ /** * Main PXF Spring Configuration class. */ -@SpringBootApplication(scanBasePackages = "org.greenplum.pxf", scanBasePackageClasses = PxfUserGroupInformation.class) +@SpringBootApplication(scanBasePackages = {"org.greenplum.pxf", "org.apache.cloudberry.pxf"}, scanBasePackageClasses = PxfUserGroupInformation.class) public class PxfServiceApplication { private static final Logger LOG = LoggerFactory.getLogger(PxfServiceApplication.class); diff --git a/server/pxf-service/src/main/java/org/greenplum/pxf/service/controller/BaseServiceImpl.java b/server/pxf-service/src/main/java/org/greenplum/pxf/service/controller/BaseServiceImpl.java index 49d25ee6c..3c59d60f8 100644 --- a/server/pxf-service/src/main/java/org/greenplum/pxf/service/controller/BaseServiceImpl.java +++ b/server/pxf-service/src/main/java/org/greenplum/pxf/service/controller/BaseServiceImpl.java @@ -1,6 +1,7 @@ package org.greenplum.pxf.service.controller; import lombok.extern.slf4j.Slf4j; +import org.apache.cloudberry.pxf.service.activity.ActiveRequestRegistry; import org.apache.hadoop.conf.Configuration; import org.greenplum.pxf.api.model.ConfigurationFactory; import org.greenplum.pxf.api.model.RequestContext; @@ -26,26 +27,30 @@ public abstract class BaseServiceImpl extends PxfErrorReporter { private final ConfigurationFactory configurationFactory; private final BridgeFactory bridgeFactory; private final SecurityService securityService; + private final ActiveRequestRegistry activeRequestRegistry; /** * Creates a new instance of the service with auto-wired dependencies. * - * @param serviceName name of the service - * @param configurationFactory configuration factory - * @param bridgeFactory bridge factory - * @param securityService security service - * @param metricsReporter metrics reporter service + * @param serviceName name of the service + * @param configurationFactory configuration factory + * @param bridgeFactory bridge factory + * @param securityService security service + * @param metricsReporter metrics reporter service + * @param activeRequestRegistry registry of in-flight requests */ protected BaseServiceImpl(String serviceName, ConfigurationFactory configurationFactory, BridgeFactory bridgeFactory, SecurityService securityService, - MetricsReporter metricsReporter) { + MetricsReporter metricsReporter, + ActiveRequestRegistry activeRequestRegistry) { this.serviceName = serviceName; this.configurationFactory = configurationFactory; this.bridgeFactory = bridgeFactory; this.securityService = securityService; this.metricsReporter = metricsReporter; + this.activeRequestRegistry = activeRequestRegistry; } /** @@ -70,8 +75,19 @@ protected OperationStats processData(RequestContext context, PrivilegedAction implements /** * Creates a new instance. * - * @param configurationFactory configuration factory - * @param bridgeFactory bridge factory - * @param securityService security service - * @param fragmenterService fragmenter service - * @param metricsReporter metrics reporter service + * @param configurationFactory configuration factory + * @param bridgeFactory bridge factory + * @param securityService security service + * @param fragmenterService fragmenter service + * @param metricsReporter metrics reporter service + * @param activeRequestRegistry registry of in-flight requests */ public ReadServiceImpl(ConfigurationFactory configurationFactory, BridgeFactory bridgeFactory, SecurityService securityService, FragmenterService fragmenterService, - MetricsReporter metricsReporter) { - super("Read", configurationFactory, bridgeFactory, securityService, metricsReporter); + MetricsReporter metricsReporter, + ActiveRequestRegistry activeRequestRegistry) { + super("Read", configurationFactory, bridgeFactory, securityService, metricsReporter, activeRequestRegistry); this.fragmenterService = fragmenterService; } @@ -84,6 +87,13 @@ private OperationResult writeStream(RequestContext context, OutputStream outputS try { List fragments = fragmenterService.getFragmentsForSegment(context); for (int i = 0; i < fragments.size(); i++) { + // stop before starting the next fragment if the request was + // cancelled via pxf_cancel_backend + if (isCancelled()) { + log.info("Read of resource {} cancelled after {} of {} fragments", + context.getDataSource(), i, fragments.size()); + break; + } Fragment fragment = fragments.get(i); sourceName = fragment.getSourceName(); String profile = fragment.getProfile(); @@ -148,6 +158,8 @@ private void processFragment(CountingOutputStream countingOutputStream, Bridge bridge = null; try { bridge = getBridge(context); + // expose the bridge so pxf_cancel_backend can end it mid-read + attachBridge(bridge); if (!bridge.beginIteration()) { log.debug("Skipping streaming fragment {} of resource {}", context.getFragmentIndex(), context.getDataSource()); @@ -167,6 +179,9 @@ private void processFragment(CountingOutputStream countingOutputStream, bridge.endIteration(); } catch (Exception e) { log.warn("Ignoring error encountered during bridge.endIteration()", e); + } finally { + // stop exposing the now-closed bridge to pxf_cancel_backend + attachBridge(null); } } Duration duration = Duration.between(startTime, Instant.now()); diff --git a/server/pxf-service/src/main/java/org/greenplum/pxf/service/controller/WriteServiceImpl.java b/server/pxf-service/src/main/java/org/greenplum/pxf/service/controller/WriteServiceImpl.java index 7b686e128..361435c81 100644 --- a/server/pxf-service/src/main/java/org/greenplum/pxf/service/controller/WriteServiceImpl.java +++ b/server/pxf-service/src/main/java/org/greenplum/pxf/service/controller/WriteServiceImpl.java @@ -2,6 +2,7 @@ import com.google.common.io.CountingInputStream; import lombok.extern.slf4j.Slf4j; +import org.apache.cloudberry.pxf.service.activity.ActiveRequestRegistry; import org.greenplum.pxf.api.model.ConfigurationFactory; import org.greenplum.pxf.api.model.RequestContext; import org.greenplum.pxf.api.utilities.Utilities; @@ -24,15 +25,18 @@ public class WriteServiceImpl extends BaseServiceImpl implements /** * Creates a new instance. * - * @param configurationFactory configuration factory - * @param bridgeFactory bridge factory - * @param securityService security service - */ + * @param configurationFactory configuration factory + * @param bridgeFactory bridge factory + * @param securityService security service + * @param metricsReporter metrics reporter service + * @param activeRequestRegistry registry of in-flight requests + * */ public WriteServiceImpl(ConfigurationFactory configurationFactory, BridgeFactory bridgeFactory, SecurityService securityService, - MetricsReporter metricsReporter) { - super("Write", configurationFactory, bridgeFactory, securityService, metricsReporter); + MetricsReporter metricsReporter, + ActiveRequestRegistry activeRequestRegistry) { + super("Write", configurationFactory, bridgeFactory, securityService, metricsReporter, activeRequestRegistry); } @Override @@ -62,11 +66,16 @@ private OperationResult readStream(RequestContext context, InputStream inputStre // dataStream (and inputStream as the result) will close automatically at the end of the try block CountingInputStream countingInputStream = new CountingInputStream(inputStream); try (DataInputStream dataStream = new DataInputStream(countingInputStream)) { + // expose the bridge so pxf_cancel_backend can end it mid-write + attachBridge(bridge); // open the output file, returns true or throws an error bridge.beginIteration(); - while (bridge.setNext(dataStream)) { + while (!isCancelled() && bridge.setNext(dataStream)) { operationStats.reportCompletedRecord(countingInputStream.getCount()); } + if (isCancelled()) { + log.info("Write to resource {} cancelled", context.getDataSource()); + } } catch (Exception e) { operationResult.setException(e); } finally { @@ -76,8 +85,12 @@ private OperationResult readStream(RequestContext context, InputStream inputStre if (operationResult.getException() == null) { operationResult.setException(e); } + } finally { + // stop exposing the now-closed bridge to pxf_cancel_backend + attachBridge(null); } + // in the case where we fail to report a record due to an exception, // report the number of bytes that we were able to read before failure operationStats.setByteCount(countingInputStream.getCount()); diff --git a/server/pxf-service/src/test/java/org/apache/cloudberry/pxf/service/activity/ActiveRequestRegistryTest.java b/server/pxf-service/src/test/java/org/apache/cloudberry/pxf/service/activity/ActiveRequestRegistryTest.java new file mode 100644 index 000000000..0dda393ad --- /dev/null +++ b/server/pxf-service/src/test/java/org/apache/cloudberry/pxf/service/activity/ActiveRequestRegistryTest.java @@ -0,0 +1,250 @@ +package org.apache.cloudberry.pxf.service.activity; + +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import org.greenplum.pxf.api.io.Writable; +import org.greenplum.pxf.api.model.RequestContext; +import org.greenplum.pxf.service.bridge.Bridge; +import org.junit.jupiter.api.BeforeEach; +import org.junit.jupiter.api.Test; + +import java.io.DataInputStream; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.atomic.AtomicInteger; + +import static org.junit.jupiter.api.Assertions.assertEquals; +import static org.junit.jupiter.api.Assertions.assertFalse; +import static org.junit.jupiter.api.Assertions.assertTrue; + +public class ActiveRequestRegistryTest { + + private ActiveRequestRegistry registry; + + @BeforeEach + public void setup() { + registry = new ActiveRequestRegistry(); + } + + private RequestContext context(int segmentId, String xid) { + return context(segmentId, xid, 0); + } + + private RequestContext context(int segmentId, String xid, int sessionId) { + RequestContext context = new RequestContext(); + context.setRequestType(RequestContext.RequestType.READ_BRIDGE); + context.setSegmentId(segmentId); + context.setTransactionId(xid); + context.setGpSessionId(sessionId); + return context; + } + + @Test + public void emptyRegistryReturnsNoActivity() { + assertTrue(registry.snapshot(ActiveRequestRegistry.ALL_SEGMENTS).isEmpty()); + assertTrue(registry.snapshot(0).isEmpty()); + } + + @Test + public void registeredRequestIsVisibleThenRemovedOnUnregister() { + registry.register(context(3, "100")); + + List all = registry.snapshot(ActiveRequestRegistry.ALL_SEGMENTS); + assertEquals(1, all.size()); + assertEquals(3, all.get(0).getSegmentId()); + assertEquals("100", all.get(0).getTransactionId()); + assertEquals("READ_BRIDGE", all.get(0).getRequestType()); + + registry.unregister(); + assertTrue(registry.snapshot(ActiveRequestRegistry.ALL_SEGMENTS).isEmpty()); + } + + @Test + public void snapshotFiltersBySegmentId() { + registry.register(context(0, "100")); + registry.register(context(1, "101")); + registry.register(context(1, "102")); + + assertEquals(3, registry.snapshot(ActiveRequestRegistry.ALL_SEGMENTS).size()); + assertEquals(1, registry.snapshot(0).size()); + assertEquals(2, registry.snapshot(1).size()); + assertTrue(registry.snapshot(2).isEmpty()); + } + + @Test + public void concurrentRequestsTrackedIndependently() throws Exception { + // Each request registers/unregisters on its own thread, as the real + // servlet workers do; entries are tracked by ActiveRequest identity. + Worker w1 = new Worker(context(0, "100"), null); + Worker w2 = new Worker(context(0, "101"), null); + w1.start(); + w2.start(); + + assertEquals(2, registry.snapshot(0).size()); + w1.finishAndJoin(); + assertEquals(1, registry.snapshot(0).size()); + w2.finishAndJoin(); + assertTrue(registry.snapshot(0).isEmpty()); + } + + @Test + public void cancelEndsBridgeAndRaisesCancelledFlagForMatchingSession() throws Exception { + CountingBridge bridge = new CountingBridge(); + Worker worker = new Worker(context(0, "100", 42), bridge); + worker.start(); + + assertEquals(1, registry.cancel(0, 42)); + assertEquals(1, bridge.endIterationCount.get()); + + worker.finishAndJoin(); + assertTrue(worker.cancelledSeen, "worker should observe the cancelled flag via isCurrentCancelled()"); + assertFalse(worker.interrupted, "cancel must not interrupt the worker thread"); + } + + @Test + public void cancelIgnoresOtherSessionsAndSegments() throws Exception { + CountingBridge bridge = new CountingBridge(); + Worker worker = new Worker(context(0, "100", 42), bridge); + worker.start(); + + assertEquals(0, registry.cancel(0, 99), "different session"); + assertEquals(0, registry.cancel(1, 42), "different segment"); + assertEquals(0, bridge.endIterationCount.get()); + + worker.finishAndJoin(); + } + + @Test + public void interruptWakesMatchingWorkerThread() throws Exception { + Worker worker = new Worker(context(0, "100", 42), null); + worker.start(); + + assertEquals(1, registry.interrupt(0, 42)); + + // interrupting unblocks the worker's interruptible wait on its own + worker.join(5000); + assertFalse(worker.isAlive(), "interrupted worker should have finished"); + assertTrue(worker.interrupted, "worker thread should have been interrupted"); + } + + @Test + public void interruptIgnoresOtherSessions() throws Exception { + Worker worker = new Worker(context(0, "100", 42), null); + worker.start(); + + assertEquals(0, registry.interrupt(0, 99)); + + worker.finishAndJoin(); + assertFalse(worker.interrupted); + } + + @Test + public void cancelAndInterruptAreNoOpsAfterUnregister() { + // register and unregister synchronously on this thread so the entry is + // marked finished; both operations must then refuse to act + registry.register(context(0, "100", 42)); + registry.unregister(); + + assertEquals(0, registry.cancel(0, 42)); + assertEquals(0, registry.interrupt(0, 42)); + assertFalse(Thread.currentThread().isInterrupted(), "the test thread must not be interrupted"); + } + + /** + * A worker that registers itself (capturing its own thread, as the real + * servlet worker does), optionally attaches a bridge, then blocks on an + * interruptible latch until released or interrupted, and finally unregisters. + */ + private final class Worker extends Thread { + private final RequestContext context; + private final Bridge bridge; + private final CountDownLatch registered = new CountDownLatch(1); + private final CountDownLatch release = new CountDownLatch(1); + volatile boolean interrupted; + volatile boolean cancelledSeen; + + Worker(RequestContext context, Bridge bridge) { + this.context = context; + this.bridge = bridge; + } + + @Override + public void run() { + registry.register(context); + if (bridge != null) { + registry.attachBridge(bridge); + } + registered.countDown(); + try { + release.await(); + } catch (InterruptedException e) { + interrupted = true; + } finally { + cancelledSeen = registry.isCurrentCancelled(); + registry.unregister(); + } + } + + @Override + public synchronized void start() { + super.start(); + try { + assertTrue(registered.await(5, TimeUnit.SECONDS), "worker failed to register"); + } catch (InterruptedException e) { + throw new RuntimeException(e); + } + } + + void finishAndJoin() throws InterruptedException { + release.countDown(); + join(5000); + assertFalse(isAlive(), "worker should have finished"); + } + } + + /** Minimal bridge that only counts endIteration() invocations. */ + private static final class CountingBridge implements Bridge { + final AtomicInteger endIterationCount = new AtomicInteger(); + final AtomicBoolean began = new AtomicBoolean(); + + @Override + public boolean beginIteration() { + began.set(true); + return true; + } + + @Override + public Writable getNext() { + return null; + } + + @Override + public boolean setNext(DataInputStream inputStream) { + return false; + } + + @Override + public void endIteration() { + endIterationCount.incrementAndGet(); + } + } +} diff --git a/server/pxf-service/src/test/java/org/apache/cloudberry/pxf/service/rest/PxfBackendControlResourceIT.java b/server/pxf-service/src/test/java/org/apache/cloudberry/pxf/service/rest/PxfBackendControlResourceIT.java new file mode 100644 index 000000000..bc4a598a3 --- /dev/null +++ b/server/pxf-service/src/test/java/org/apache/cloudberry/pxf/service/rest/PxfBackendControlResourceIT.java @@ -0,0 +1,95 @@ +package org.apache.cloudberry.pxf.service.rest; + +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import org.greenplum.pxf.service.HttpHeaderDecoder; +import org.greenplum.pxf.service.PxfServiceApplication; +import org.apache.cloudberry.pxf.service.activity.ActiveRequestRegistry; +import org.junit.jupiter.api.Test; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.test.autoconfigure.web.servlet.WebMvcTest; +import org.springframework.boot.test.mock.mockito.MockBean; +import org.springframework.test.context.ContextConfiguration; +import org.springframework.test.web.servlet.MockMvc; + +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; +import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.jsonPath; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status; + +@WebMvcTest(PxfBackendControlResource.class) +@ContextConfiguration(classes = PxfServiceApplication.class) +public class PxfBackendControlResourceIT { + + @Autowired + private MockMvc mvc; + + @MockBean + private ActiveRequestRegistry mockRegistry; + + // required by web-layer beans pulled into the @WebMvcTest slice + @MockBean + private HttpHeaderDecoder mockHttpHeaderDecoder; + + @Test + public void cancelBackendFiltersBySegmentAndSessionHeaders() throws Exception { + when(mockRegistry.cancel(eq(3), eq(42))).thenReturn(2); + + mvc.perform(get("/pxf/cancel_backend") + .header("X-GP-SEGMENT-ID", "3") + .header("X-GP-SESSION-ID", "42")) + .andExpect(status().isOk()) + .andExpect(jsonPath("$.cancelled").value(2)); + + verify(mockRegistry).cancel(3, 42); + } + + @Test + public void cancelBackendWithoutSegmentHeaderTargetsAllSegments() throws Exception { + when(mockRegistry.cancel(eq(ActiveRequestRegistry.ALL_SEGMENTS), eq(42))).thenReturn(1); + + mvc.perform(get("/pxf/cancel_backend").header("X-GP-SESSION-ID", "42")) + .andExpect(status().isOk()) + .andExpect(jsonPath("$.cancelled").value(1)); + + verify(mockRegistry).cancel(ActiveRequestRegistry.ALL_SEGMENTS, 42); + } + + @Test + public void interruptBackendFiltersBySegmentAndSessionHeaders() throws Exception { + when(mockRegistry.interrupt(eq(3), eq(42))).thenReturn(1); + + mvc.perform(get("/pxf/interrupt_backend") + .header("X-GP-SEGMENT-ID", "3") + .header("X-GP-SESSION-ID", "42")) + .andExpect(status().isOk()) + .andExpect(jsonPath("$.interrupted").value(1)); + + verify(mockRegistry).interrupt(3, 42); + } + + @Test + public void missingSessionHeaderIsBadRequest() throws Exception { + mvc.perform(get("/pxf/cancel_backend").header("X-GP-SEGMENT-ID", "3")) + .andExpect(status().isBadRequest()); + } +} diff --git a/server/pxf-service/src/test/java/org/apache/cloudberry/pxf/service/rest/PxfStatActivityResourceIT.java b/server/pxf-service/src/test/java/org/apache/cloudberry/pxf/service/rest/PxfStatActivityResourceIT.java new file mode 100644 index 000000000..18d89267d --- /dev/null +++ b/server/pxf-service/src/test/java/org/apache/cloudberry/pxf/service/rest/PxfStatActivityResourceIT.java @@ -0,0 +1,92 @@ +package org.apache.cloudberry.pxf.service.rest; + +/* + * Licensed to the Apache Software Foundation (ASF) under one + * or more contributor license agreements. See the NOTICE file + * distributed with this work for additional information + * regarding copyright ownership. The ASF licenses this file + * to you under the Apache License, Version 2.0 (the + * "License"); you may not use this file except in compliance + * with the License. You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, + * software distributed under the License is distributed on an + * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY + * KIND, either express or implied. See the License for the + * specific language governing permissions and limitations + * under the License. + */ + +import org.greenplum.pxf.api.model.RequestContext; +import org.greenplum.pxf.service.HttpHeaderDecoder; +import org.greenplum.pxf.service.PxfServiceApplication; +import org.apache.cloudberry.pxf.service.activity.ActiveRequestInfo; +import org.apache.cloudberry.pxf.service.activity.ActiveRequestRegistry; +import org.junit.jupiter.api.Test; +import org.springframework.beans.factory.annotation.Autowired; +import org.springframework.boot.test.autoconfigure.web.servlet.WebMvcTest; +import org.springframework.boot.test.mock.mockito.MockBean; +import org.springframework.test.context.ContextConfiguration; +import org.springframework.test.web.servlet.MockMvc; + +import java.util.Collections; + +import static org.mockito.ArgumentMatchers.eq; +import static org.mockito.Mockito.when; +import static org.springframework.test.web.servlet.request.MockMvcRequestBuilders.get; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.jsonPath; +import static org.springframework.test.web.servlet.result.MockMvcResultMatchers.status; + +@WebMvcTest(PxfStatActivityResource.class) +@ContextConfiguration(classes = PxfServiceApplication.class) +public class PxfStatActivityResourceIT { + + @Autowired + private MockMvc mvc; + + @MockBean + private ActiveRequestRegistry mockRegistry; + + @MockBean + private HttpHeaderDecoder mockHttpHeaderDecoder; + + private static ActiveRequestInfo sampleInfo(int segmentId, String xid) { + RequestContext context = new RequestContext(); + context.setRequestType(RequestContext.RequestType.READ_BRIDGE); + context.setSegmentId(segmentId); + context.setTransactionId(xid); + context.setGpSessionId(42); + context.setGpCommandCount(7); + context.setProfile("hdfs:text"); + return new ActiveRequestInfo(context, 1718700000000L, "sdw3"); + } + + @Test + public void emptyRegistryReturnsEmptyActivities() throws Exception { + when(mockRegistry.snapshot(ActiveRequestRegistry.ALL_SEGMENTS)) + .thenReturn(Collections.emptyList()); + + mvc.perform(get("/pxf/stat_activity")) + .andExpect(status().isOk()) + .andExpect(jsonPath("$.activities").isArray()) + .andExpect(jsonPath("$.activities").isEmpty()); + } + + @Test + public void filtersBySegmentIdHeaderAndSerializesFields() throws Exception { + when(mockRegistry.snapshot(eq(3))) + .thenReturn(Collections.singletonList(sampleInfo(3, "1234"))); + + mvc.perform(get("/pxf/stat_activity").header("X-GP-SEGMENT-ID", "3")) + .andExpect(status().isOk()) + .andExpect(jsonPath("$.activities[0].segmentId").value(3)) + .andExpect(jsonPath("$.activities[0].transactionId").value("1234")) + .andExpect(jsonPath("$.activities[0].requestType").value("READ_BRIDGE")) + .andExpect(jsonPath("$.activities[0].gpSessionId").value(42)) + .andExpect(jsonPath("$.activities[0].profile").value("hdfs:text")) + .andExpect(jsonPath("$.activities[0].host").value("sdw3")) + .andExpect(jsonPath("$.activities[0].startTimeMs").value(1718700000000L)); + } +} diff --git a/server/pxf-service/src/test/java/org/greenplum/pxf/service/controller/ReadServiceImplTest.java b/server/pxf-service/src/test/java/org/greenplum/pxf/service/controller/ReadServiceImplTest.java index ed03a3a07..a2db88b17 100644 --- a/server/pxf-service/src/test/java/org/greenplum/pxf/service/controller/ReadServiceImplTest.java +++ b/server/pxf-service/src/test/java/org/greenplum/pxf/service/controller/ReadServiceImplTest.java @@ -1,5 +1,6 @@ package org.greenplum.pxf.service.controller; +import org.apache.cloudberry.pxf.service.activity.ActiveRequestRegistry; import org.apache.hadoop.conf.Configuration; import org.greenplum.pxf.api.error.PxfRuntimeException; import org.greenplum.pxf.api.io.Writable; @@ -75,7 +76,7 @@ public void setup() throws Exception { return result; }); - readService = new ReadServiceImpl(mockConfigurationFactory, mockBridgeFactory, mockSecurityService, mockFragmenterService, mockMetricReporter); + readService = new ReadServiceImpl(mockConfigurationFactory, mockBridgeFactory, mockSecurityService, mockFragmenterService, mockMetricReporter, new ActiveRequestRegistry()); } @Test diff --git a/server/pxf-service/src/test/java/org/greenplum/pxf/service/controller/WriteServiceImplTest.java b/server/pxf-service/src/test/java/org/greenplum/pxf/service/controller/WriteServiceImplTest.java index 784fe940b..c0c617eae 100644 --- a/server/pxf-service/src/test/java/org/greenplum/pxf/service/controller/WriteServiceImplTest.java +++ b/server/pxf-service/src/test/java/org/greenplum/pxf/service/controller/WriteServiceImplTest.java @@ -1,5 +1,6 @@ package org.greenplum.pxf.service.controller; +import org.apache.cloudberry.pxf.service.activity.ActiveRequestRegistry; import org.apache.hadoop.conf.Configuration; import org.greenplum.pxf.api.model.ConfigurationFactory; import org.greenplum.pxf.api.model.RequestContext; @@ -61,7 +62,7 @@ public void setup() throws Exception { }); when(mockBridgeFactory.getBridge(mockContext)).thenReturn(mockBridge); - writeService = new WriteServiceImpl(mockConfigurationFactory, mockBridgeFactory, mockSecurityService, mockMetricReporter); + writeService = new WriteServiceImpl(mockConfigurationFactory, mockBridgeFactory, mockSecurityService, mockMetricReporter, new ActiveRequestRegistry()); } @Test