Fix issues around EXPLAIN with JIT.

I (Andres) was more than a bit hasty in committing 33001fd7a7072d48327
after last minute changes, leading to a number of problems (jit output
was only shown for JIT in parallel workers, and just EXPLAIN without
ANALYZE didn't work).  Lukas luckily found these issues quickly.

Instead of combining instrumentation in in standard_ExecutorEnd(), do
so on demand in the new ExplainPrintJITSummary().

Also update a documentation example of the JIT output, changed in
52050ad8ebec8d831.

Author: Lukas Fittl, with minor changes by me
Discussion: https://postgr.es/m/CAP53PkxmgJht69pabxBXJBM+0oc6kf3KHMborLP7H2ouJ0CCtQ@mail.gmail.com
Backpatch: 11, where JIT compilation was introduced

7 files changed, 39 insertions, 34 deletions

diff --git a/contrib/auto_explain/auto_explain.c b/contrib/auto_explain/auto_explain.c
index a19c1d2a3c3..646cd0d42cd 100644
--- a/contrib/auto_explain/auto_explain.c
+++ b/contrib/auto_explain/auto_explain.c
@@ -363,8 +363,7 @@ explain_ExecutorEnd(QueryDesc *queryDesc)
 			if (es->analyze && auto_explain_log_triggers)
 				ExplainPrintTriggers(es, queryDesc);
 			if (es->costs)
-				ExplainPrintJIT(es, queryDesc->estate->es_jit_flags,
-								queryDesc->estate->es_jit_combined_instr, -1);
+				ExplainPrintJITSummary(es, queryDesc);
 			ExplainEndOutput(es);
 
 			/* Remove last line break */
diff --git a/doc/src/sgml/jit.sgml b/doc/src/sgml/jit.sgml
index 8387a4f6bee..a21a07ef71d 100644
--- a/doc/src/sgml/jit.sgml
+++ b/doc/src/sgml/jit.sgml
@@ -169,12 +169,8 @@ SET
  Planning Time: 0.133 ms
  JIT:
    Functions: 3
-   Generation Time: 1.259 ms
-   Inlining: false
-   Inlining Time: 0.000 ms
-   Optimization: false
-   Optimization Time: 0.797 ms
-   Emission Time: 5.048 ms
+   Options: Inlining false, Optimization false, Expressions true, Deforming true
+   Timing: Generation 1.259 ms, Inlining 0.000 ms, Optimization 0.797 ms, Emission 5.048 ms, Total 7.104 ms
  Execution Time: 7.416 ms
 </screen>
    As visible here, <acronym>JIT</acronym> was used, but inlining and
diff --git a/src/backend/commands/explain.c b/src/backend/commands/explain.c
index ed6afe79a9f..799a22e9d55 100644
--- a/src/backend/commands/explain.c
+++ b/src/backend/commands/explain.c
@@ -564,8 +564,7 @@ ExplainOnePlan(PlannedStmt *plannedstmt, IntoClause *into, ExplainState *es,
 	 * at a later stage.
 	 */
 	if (es->costs)
-		ExplainPrintJIT(es, queryDesc->estate->es_jit_flags,
-						queryDesc->estate->es_jit_combined_instr, -1);
+		ExplainPrintJITSummary(es, queryDesc);
 
 	/*
 	 * Close down the query and free resources.  Include time for this in the
@@ -689,6 +688,32 @@ ExplainPrintTriggers(ExplainState *es, QueryDesc *queryDesc)
 }
 
 /*
+ * ExplainPrintJITSummary -
+ *    Print summarized JIT instrumentation from leader and workers
+ */
+void
+ExplainPrintJITSummary(ExplainState *es, QueryDesc *queryDesc)
+{
+	JitInstrumentation ji = {0};
+
+	if (!(queryDesc->estate->es_jit_flags & PGJIT_PERFORM))
+		return;
+
+	/*
+	 * Work with a copy instead of modifying the leader state, since this
+	 * function may be called twice
+	 */
+	if (queryDesc->estate->es_jit)
+		InstrJitAgg(&ji, &queryDesc->estate->es_jit->instr);
+
+	/* If this process has done JIT in parallel workers, merge stats */
+	if (queryDesc->estate->es_jit_worker_instr)
+		InstrJitAgg(&ji, queryDesc->estate->es_jit_worker_instr);
+
+	ExplainPrintJIT(es, queryDesc->estate->es_jit_flags, &ji, -1);
+}
+
+/*
  * ExplainPrintJIT -
  *	  Append information about JITing to es->str.
  *
diff --git a/src/backend/executor/execMain.c b/src/backend/executor/execMain.c
index c865032a80f..1e97ff5643f 100644
--- a/src/backend/executor/execMain.c
+++ b/src/backend/executor/execMain.c
@@ -495,21 +495,6 @@ standard_ExecutorEnd(QueryDesc *queryDesc)
 
 	ExecEndPlan(queryDesc->planstate, estate);
 
-	/*
-	 * If this process has done JIT, either merge stats into worker stats, or
-	 * use this process' stats as the global stats if no parallelism was used
-	 * / no workers did JIT.
-	 */
-	if (estate->es_instrument && queryDesc->estate->es_jit)
-	{
-		if (queryDesc->estate->es_jit_combined_instr)
-			InstrJitAgg(queryDesc->estate->es_jit_combined_instr,
-						&queryDesc->estate->es_jit->instr);
-		else
-			queryDesc->estate->es_jit_combined_instr =
-				&queryDesc->estate->es_jit->instr;
-	}
-
 	/* do away with our snapshots */
 	UnregisterSnapshot(estate->es_snapshot);
 	UnregisterSnapshot(estate->es_crosscheck_snapshot);
diff --git a/src/backend/executor/execParallel.c b/src/backend/executor/execParallel.c
index 7d8bd01994f..0fdbd119d9b 100644
--- a/src/backend/executor/execParallel.c
+++ b/src/backend/executor/execParallel.c
@@ -1043,13 +1043,12 @@ ExecParallelRetrieveJitInstrumentation(PlanState *planstate,
 
 	/*
 	 * Accumulate worker JIT instrumentation into the combined JIT
-	 * instrumentation, allocating it if required. Note this is kept separate
-	 * from the leader's own instrumentation.
+	 * instrumentation, allocating it if required.
 	 */
-	if (!planstate->state->es_jit_combined_instr)
-		planstate->state->es_jit_combined_instr =
+	if (!planstate->state->es_jit_worker_instr)
+		planstate->state->es_jit_worker_instr =
 			MemoryContextAllocZero(planstate->state->es_query_cxt, sizeof(JitInstrumentation));
-	combined = planstate->state->es_jit_combined_instr;
+	combined = planstate->state->es_jit_worker_instr;
 
 	/* Accummulate all the workers' instrumentations. */
 	for (n = 0; n < shared_jit->num_workers; ++n)
diff --git a/src/include/commands/explain.h b/src/include/commands/explain.h
index b0b6f1e15a4..d3f70fda084 100644
--- a/src/include/commands/explain.h
+++ b/src/include/commands/explain.h
@@ -81,6 +81,7 @@ extern void ExplainOnePlan(PlannedStmt *plannedstmt, IntoClause *into,
 extern void ExplainPrintPlan(ExplainState *es, QueryDesc *queryDesc);
 extern void ExplainPrintTriggers(ExplainState *es, QueryDesc *queryDesc);
 
+extern void ExplainPrintJITSummary(ExplainState *es, QueryDesc *queryDesc);
 extern void ExplainPrintJIT(ExplainState *es, int jit_flags,
 				struct JitInstrumentation *jit_instr, int worker_i);
 
diff --git a/src/include/nodes/execnodes.h b/src/include/nodes/execnodes.h
index 03ad5169762..020ecdcd404 100644
--- a/src/include/nodes/execnodes.h
+++ b/src/include/nodes/execnodes.h
@@ -570,13 +570,13 @@ typedef struct EState
 	 * and with which options.  es_jit is created on-demand when JITing is
 	 * performed.
 	 *
-	 * es_jit_combined_instr, at the end of query execution with
-	 * instrumentation enabled, is the the combined instrumentation
-	 * information of leader and followers.
+	 * es_jit_combined_instr is the the combined, on demand allocated,
+	 * instrumentation from all workers. The leader's instrumentation is kept
+	 * separate, and is combined on demand by ExplainPrintJITSummary().
 	 */
 	int			es_jit_flags;
 	struct JitContext *es_jit;
-	struct JitInstrumentation *es_jit_combined_instr;
+	struct JitInstrumentation *es_jit_worker_instr;
 } EState;