Skip to content
4 changes: 4 additions & 0 deletions src/backend/tcop/postgres.c
Original file line number Diff line number Diff line change
Expand Up @@ -107,6 +107,10 @@ int client_connection_check_interval = 0;
/* flags for non-system relation kinds to restrict use */
int restrict_nonsystem_relation_kind;

#ifdef USE_INJECTION_POINTS
/* Counter which can show us if we are in SMGR */
int inside_smgr_api = 0;
#endif
/* ----------------
* private typedefs etc
* ----------------
Expand Down
16 changes: 16 additions & 0 deletions src/include/miscadmin.h
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,10 @@
#include "datatype/timestamp.h" /* for TimestampTz */
#include "pgtime.h" /* for pg_time_t */

#ifdef USE_INJECTION_POINTS
#include "utils/injection_point.h"
#endif


#define InvalidPid (-1)

Expand Down Expand Up @@ -123,9 +127,21 @@ extern process_interrupts_callback_t ProcessInterruptsCallback;
unlikely(InterruptPending))
#endif

#ifdef USE_INJECTION_POINTS
/* Run injection point when inside smgr API; used by CHECK_FOR_INTERRUPTS */
#define CHECK_FOR_INTERRUPTS_SMGR_INJECTION() \
do { \
if (INTERRUPTS_CAN_BE_PROCESSED() && inside_smgr_api > 0) \
INJECTION_POINT("SMGR_API"); \
} while(0)
#else
#define CHECK_FOR_INTERRUPTS_SMGR_INJECTION() ((void)0)
#endif

/* Service interrupt, if one is pending and it's safe to service it now */
#define CHECK_FOR_INTERRUPTS() \
do { \
CHECK_FOR_INTERRUPTS_SMGR_INJECTION(); \
if (INTERRUPTS_PENDING_CONDITION()) \
ProcessInterrupts(); \
} while(0)
Expand Down
5 changes: 5 additions & 0 deletions src/include/postgres.h
Original file line number Diff line number Diff line change
Expand Up @@ -577,3 +577,8 @@ extern Datum Float8GetDatum(float8 X);
#endif

#endif /* POSTGRES_H */

/* Declare for neon use */
#ifdef USE_INJECTION_POINTS
extern int inside_smgr_api;
#endif
50 changes: 50 additions & 0 deletions src/test/modules/injection_points/injection_points.c
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
*/

#include "postgres.h"
#include <time.h>

#include "fmgr.h"
#include "miscadmin.h"
Expand Down Expand Up @@ -59,6 +60,8 @@ typedef struct InjectionPointCondition

/* ID of the process where the injection point is allowed to run */
int pid;
/* probability in [0,1], 1.0 = always */
double prob;
} InjectionPointCondition;

/*
Expand Down Expand Up @@ -88,13 +91,16 @@ static InjectionPointSharedState *inj_state = NULL;

extern PGDLLEXPORT void injection_error(const char *name,
const void *private_data);
extern PGDLLEXPORT void injection_error_prob(const char *name,
const void *private_data);
extern PGDLLEXPORT void injection_notice(const char *name,
const void *private_data);
extern PGDLLEXPORT void injection_wait(const char *name,
const void *private_data);

/* track if injection points attached in this process are linked to it */
static bool injection_point_local = false;
static double action2prob(const char *action, int pos);

/*
* Callback for shared memory area initialization.
Expand Down Expand Up @@ -184,7 +190,20 @@ injection_error(const char *name, const void *private_data)

elog(ERROR, "error triggered for injection point %s", name);
}
void
injection_error_prob(const char *name, const void *private_data)
{
InjectionPointCondition *condition = (InjectionPointCondition *) private_data;

if (!injection_point_allowed(condition))
return;

/* Use the probability stored in the condition. */
if ((double) rand() / (double) RAND_MAX > condition->prob)
return;

elog(ERROR, "error triggered for injection point %s", name);
}
void
injection_notice(const char *name, const void *private_data)
{
Expand Down Expand Up @@ -278,6 +297,11 @@ injection_points_attach(PG_FUNCTION_ARGS)
function = "injection_notice";
else if (strcmp(action, "wait") == 0)
function = "injection_wait";
else if (strncmp(action, "error-prob-", 11) == 0)
{
condition.prob = action2prob(action, 11);
function = "injection_error_prob";
}
else
elog(ERROR, "incorrect action \"%s\" for injection point creation", action);

Expand Down Expand Up @@ -402,3 +426,29 @@ injection_points_detach(PG_FUNCTION_ARGS)

PG_RETURN_VOID();
}

/*
* Coverts the action name into probability
*/
static double action2prob(const char *action, const int pos)
{
/*
* Simple parser: convert "0-01" -> "0.01" then strtod().
*/
const char *p = action + pos; /* points to "0-01" */
double prob;
char *endptr;
char buf[32];
int i, j;

for (i = 0, j = 0; p[i] != '\0' && j < (int) sizeof(buf) - 1; i++)
{
buf[j++] = (p[i] == '-') ? '.' : p[i];
}
buf[j] = '\0';
errno = 0;
prob = strtod(buf, &endptr);
if (errno != 0 || endptr == buf || prob < 0.0 || prob > 1.0)
elog(ERROR, "invalid probability in action \"%s\"", action);
return prob;
}