mirror of https://github.com/citusdata/citus.git
152 lines
3.4 KiB
C
152 lines
3.4 KiB
C
/*-------------------------------------------------------------------------
|
|
*
|
|
* progress_utils.c
|
|
*
|
|
* This file contains functions to exercise progress monitoring functionality
|
|
* within Citus.
|
|
*
|
|
* Copyright (c) 2017, Citus Data, Inc.
|
|
*
|
|
*-------------------------------------------------------------------------
|
|
*/
|
|
|
|
|
|
#include "postgres.h"
|
|
#include "miscadmin.h"
|
|
#include "fmgr.h"
|
|
#include "funcapi.h"
|
|
|
|
#include <unistd.h>
|
|
|
|
#include "distributed/multi_progress.h"
|
|
#include "nodes/execnodes.h"
|
|
#include "utils/tuplestore.h"
|
|
|
|
|
|
PG_FUNCTION_INFO_V1(create_progress);
|
|
PG_FUNCTION_INFO_V1(update_progress);
|
|
PG_FUNCTION_INFO_V1(finish_progress);
|
|
PG_FUNCTION_INFO_V1(show_progress);
|
|
|
|
|
|
Datum
|
|
create_progress(PG_FUNCTION_ARGS)
|
|
{
|
|
uint64 magicNumber = PG_GETARG_INT64(0);
|
|
int stepCount = PG_GETARG_INT32(1);
|
|
ProgressMonitorData *monitor = CreateProgressMonitor(magicNumber, stepCount,
|
|
sizeof(uint64), 0);
|
|
|
|
if (monitor != NULL)
|
|
{
|
|
uint64 *steps = (uint64 *) monitor->steps;
|
|
|
|
int i = 0;
|
|
for (; i < stepCount; i++)
|
|
{
|
|
steps[i] = 0;
|
|
}
|
|
}
|
|
|
|
PG_RETURN_VOID();
|
|
}
|
|
|
|
|
|
Datum
|
|
update_progress(PG_FUNCTION_ARGS)
|
|
{
|
|
uint64 step = PG_GETARG_INT64(0);
|
|
uint64 newValue = PG_GETARG_INT64(1);
|
|
|
|
ProgressMonitorData *monitor = GetCurrentProgressMonitor();
|
|
|
|
if (monitor != NULL && step < monitor->stepCount)
|
|
{
|
|
uint64 *steps = (uint64 *) monitor->steps;
|
|
steps[step] = newValue;
|
|
}
|
|
|
|
PG_RETURN_VOID();
|
|
}
|
|
|
|
|
|
Datum
|
|
finish_progress(PG_FUNCTION_ARGS)
|
|
{
|
|
FinalizeCurrentProgressMonitor();
|
|
|
|
PG_RETURN_VOID();
|
|
}
|
|
|
|
|
|
Datum
|
|
show_progress(PG_FUNCTION_ARGS)
|
|
{
|
|
uint64 magicNumber = PG_GETARG_INT64(0);
|
|
List *attachedDSMSegments = NIL;
|
|
List *monitorList = ProgressMonitorList(magicNumber, &attachedDSMSegments);
|
|
Tuplestorestate *tupstore = NULL;
|
|
TupleDesc tupdesc;
|
|
MemoryContext perQueryContext;
|
|
MemoryContext currentContext;
|
|
ReturnSetInfo *resultSet = (ReturnSetInfo *) fcinfo->resultinfo;
|
|
ListCell *monitorCell = NULL;
|
|
|
|
/* check to see if caller supports us returning a tuplestore */
|
|
if (resultSet == NULL || !IsA(resultSet, ReturnSetInfo))
|
|
{
|
|
ereport(ERROR,
|
|
(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
|
|
errmsg("set-valued function called in context that cannot " \
|
|
"accept a set")));
|
|
}
|
|
if (!(resultSet->allowedModes & SFRM_Materialize))
|
|
{
|
|
ereport(ERROR,
|
|
(errcode(ERRCODE_FEATURE_NOT_SUPPORTED),
|
|
errmsg("materialize mode required, but it is not " \
|
|
"allowed in this context")));
|
|
}
|
|
if (get_call_result_type(fcinfo, NULL, &tupdesc) != TYPEFUNC_COMPOSITE)
|
|
{
|
|
elog(ERROR, "return type must be a row type");
|
|
}
|
|
|
|
perQueryContext = resultSet->econtext->ecxt_per_query_memory;
|
|
currentContext = MemoryContextSwitchTo(perQueryContext);
|
|
tupstore = tuplestore_begin_heap(true, false, work_mem);
|
|
resultSet->returnMode = SFRM_Materialize;
|
|
resultSet->setResult = tupstore;
|
|
resultSet->setDesc = tupdesc;
|
|
MemoryContextSwitchTo(currentContext);
|
|
|
|
foreach(monitorCell, monitorList)
|
|
{
|
|
ProgressMonitorData *monitor = lfirst(monitorCell);
|
|
uint64 *steps = monitor->steps;
|
|
|
|
int stepIndex = 0;
|
|
for (stepIndex = 0; stepIndex < monitor->stepCount; stepIndex++)
|
|
{
|
|
uint64 step = steps[stepIndex];
|
|
|
|
Datum values[2];
|
|
bool nulls[2];
|
|
|
|
memset(values, 0, sizeof(values));
|
|
memset(nulls, 0, sizeof(nulls));
|
|
|
|
values[0] = Int32GetDatum(stepIndex);
|
|
values[1] = UInt64GetDatum(step);
|
|
|
|
tuplestore_putvalues(tupstore, tupdesc, values, nulls);
|
|
}
|
|
}
|
|
|
|
tuplestore_donestoring(tupstore);
|
|
|
|
DetachFromDSMSegments(attachedDSMSegments);
|
|
|
|
return (Datum) 0;
|
|
}
|