[PATCH v2 09/19] Use usdt_prids map to call clauses conditionally for USDT probes

eugene.loh
2024-09-08T16:40:26Z

This version supports only up to 64 clauses for an underlying
probe, but it can be extended to more clauses.

Signed-off-by: Eugene Loh <eugene.loh@oracle.com>
---
libdtrace/dt_prov_uprobe.c | 159 +++++++++++++++++++++++++++----------
1 file changed, 116 insertions(+), 43 deletions(-)

diff --git a/libdtrace/dt_prov_uprobe.c b/libdtrace/dt_prov_uprobe.c
index bc929878..af188061 100644
--- a/libdtrace/dt_prov_uprobe.c
+++ b/libdtrace/dt_prov_uprobe.c
@@ -681,11 +681,16 @@ static void enable_usdt(dtrace_hdl_t *dtp, dt_probe_t *prp)
*/
static int trampoline(dt_pcb_t *pcb, uint_t exitlbl)
{
+ dtrace_hdl_t *dtp = pcb->pcb_hdl;
dt_irlist_t *dlp = &pcb->pcb_ir;
const dt_probe_t *uprp = pcb->pcb_probe;
const dt_uprobe_t *upp = uprp->prv_data;
const list_probe_t *pop;
uint_t lbl_exit = pcb->pcb_exitlbl;
+ dt_ident_t *usdt_prids = dt_dlib_get_map(dtp, "usdt_prids");
+ int n;
+
+ assert(usdt_prids != NULL);

dt_cg_tramp_prologue(pcb); // call this only once... is PRID set/relocated correctly?

@@ -697,14 +702,17 @@ static int trampoline(dt_pcb_t *pcb, uint_t exitlbl)
dt_cg_tramp_copy_regs(pcb); // call this only once for all clauses?

/*
- * Loop over overlying probes, calling clauses for those that match:
+ * pid probes.
+ *
+ * Loop over overlying pid probes, calling clauses for those that match:
*
- * for overlying probes (that match except possibly for pid)
+ * for overlying pid probes (that match except possibly for pid)
* if (pid matches) {
* dctx->mst->prid = PRID1;
* < any number of clause calls >
* }
*/
+
for (pop = dt_list_next(&upp->probes); pop != NULL;
pop = dt_list_next(pop)) {
const dt_probe_t *prp = pop->probe;
@@ -712,6 +720,9 @@ static int trampoline(dt_pcb_t *pcb, uint_t exitlbl)
pid_t pid;
dt_ident_t *idp;

+ if (prp->prov->impl != &dt_pid)
+ continue;
+
pid = dt_pid_get_pid(prp->desc, pcb->pcb_hdl, pcb, NULL);
assert(pid != -1);

@@ -743,6 +754,86 @@ static int trampoline(dt_pcb_t *pcb, uint_t exitlbl)
BPF_NOP());
}

+ /*
+ * USDT
+ */
+
+ /* In some cases, we know there are no USDT probes. */ // FIXME: add more checks
+ if (upp->flags & PP_IS_RETURN)
+ goto out;
+
+ dt_cg_tramp_copy_args_from_regs(pcb, 0);
+
+ /*
+ * Retrieve the PID of the process that caused the probe to fire.
+ */
+ emit(dlp, BPF_CALL_HELPER(BPF_FUNC_get_current_pid_tgid));
+ emit(dlp, BPF_ALU64_IMM(BPF_RSH, BPF_REG_0, 32));
+
+ /*
+ * Look up in the BPF 'usdt_prids' map. Space for the look-up key
+ * will be used on the BPF stack:
+ *
+ * offset value
+ *
+ * -sizeof(usdt_prids_map_key_t) pid (in %r0)
+ *
+ * -sizeof(usdt_prids_map_key_t) + sizeof(pid_t)
+ * ==
+ * -sizeof(dtrace_id_t) underlying-probe prid
+ */
+ emit(dlp, BPF_STORE(BPF_W, BPF_REG_9, (int)(-sizeof(usdt_prids_map_key_t)), BPF_REG_0));
+ emit(dlp, BPF_STORE_IMM(BPF_W, BPF_REG_9, (int)(-sizeof(dtrace_id_t)), uprp->desc->id));
+ dt_cg_xsetx(dlp, usdt_prids, DT_LBL_NONE, BPF_REG_1, usdt_prids->di_id);
+ emit(dlp, BPF_MOV_REG(BPF_REG_2, BPF_REG_9));
+ emit(dlp, BPF_ALU64_IMM(BPF_ADD, BPF_REG_2, (int)(-sizeof(usdt_prids_map_key_t))));
+ emit(dlp, BPF_CALL_HELPER(BPF_FUNC_map_lookup_elem));
+ emit(dlp, BPF_BRANCH_IMM(BPF_JEQ, BPF_REG_0, 0, lbl_exit));
+
+ /* Read the PRID from the table lookup and store to mst->prid. */
+ emit(dlp, BPF_LOAD(BPF_W, BPF_REG_1, BPF_REG_0, 0));
+ emit(dlp, BPF_STORE(BPF_W, BPF_REG_7, DMST_PRID, BPF_REG_1));
+
+ /* Read the bit mask from the table lookup in %r6. */ // FIXME someday, extend this past 64 bits
+ emit(dlp, BPF_LOAD(BPF_DW, BPF_REG_6, BPF_REG_0, offsetof(usdt_prids_map_val_t, mask)));
+
+ /*
+ * Hold the bit mask in %r6 between clause calls.
+ */
+ for (n = 0; n < dtp->dt_stmt_nextid; n++) {
+ dt_ident_t *idp = dtp->dt_stmts[n]->dtsd_clause;
+ uint_t lbl_next = dt_irlist_label(dlp);
+
+ /* If the lowest %r6 bit is 0, skip over this clause. */
+ emit(dlp, BPF_MOV_REG(BPF_REG_1, BPF_REG_6));
+ emit(dlp, BPF_ALU64_IMM(BPF_AND, BPF_REG_1, 1));
+ emit(dlp, BPF_BRANCH_IMM(BPF_JEQ, BPF_REG_1, 0, lbl_next));
+
+ /*
+ * if (*dctx.act != act) // ldw %r0, [%r9 + DCTX_ACT]
+ * goto exit; // ldw %r0, [%r0 + 0]
+ * // jne %r0, act, lbl_exit
+ */
+ emit(dlp, BPF_LOAD(BPF_DW, BPF_REG_0, BPF_REG_9, DCTX_ACT));
+ emit(dlp, BPF_LOAD(BPF_W, BPF_REG_0, BPF_REG_0, 0));
+ emit(dlp, BPF_BRANCH_IMM(BPF_JNE, BPF_REG_0, DT_ACTIVITY_ACTIVE, lbl_exit));
+
+ /* dctx.mst->scratch_top = 8 */
+ emit(dlp, BPF_STORE_IMM(BPF_W, BPF_REG_7, DMST_SCRATCH_TOP, 8));
+
+ /* Call clause. */
+ emit(dlp, BPF_MOV_REG(BPF_REG_1, BPF_REG_9));
+ emite(dlp, BPF_CALL_FUNC(idp->di_id), idp);
+
+ /* Finished this clause. */
+ emitl(dlp, lbl_next,
+ BPF_NOP());
+
+ /* Right-shift %r6. */
+ emit(dlp, BPF_ALU64_IMM(BPF_RSH, BPF_REG_6, 1));
+ }
+
+out:
dt_cg_tramp_return(pcb);

return 0;
@@ -788,10 +879,9 @@ static int trampoline_is_enabled(dt_pcb_t *pcb, uint_t exitlbl)
{
dt_irlist_t *dlp = &pcb->pcb_ir;
const dt_probe_t *uprp = pcb->pcb_probe;
- const dt_uprobe_t *upp = uprp->prv_data;
- const list_probe_t *pop;
- uint_t lbl_assign = dt_irlist_label(dlp);
- uint_t lbl_exit = pcb->pcb_exitlbl;
+ dt_ident_t *usdt_prids = dt_dlib_get_map(pcb->pcb_hdl, "usdt_prids");
+
+ assert(usdt_prids != NULL);

dt_cg_tramp_prologue(pcb);

@@ -800,7 +890,6 @@ static int trampoline_is_enabled(dt_pcb_t *pcb, uint_t exitlbl)
* // (%r7 = dctx->mst)
* // (%r8 = dctx->ctx)
*/
-
dt_cg_tramp_copy_regs(pcb);

/*
@@ -818,46 +907,30 @@ static int trampoline_is_enabled(dt_pcb_t *pcb, uint_t exitlbl)
emit(dlp, BPF_ALU64_IMM(BPF_RSH, BPF_REG_0, 32));

/*
- * Generate a composite conditional clause, as above, except that rather
- * than emitting call_clauses, we emit copyouts instead, using
- * copyout_val() above:
+ * Look up in the BPF 'usdt_prids' map. Space for the look-up key
+ * will be used on the BPF stack:
*
- * if (pid == PID1) {
- * goto assign;
- * } else if (pid == PID2) {
- * goto assign;
- * } else if (pid == ...) {
- * goto assign;
- * }
- * goto exit;
- * assign:
- * *arg0 = 1;
- * goto exit;
+ * offset value
+ *
+ * -sizeof(usdt_prids_map_key_t) pid (in %r0)
*
- * It is valid and safe to use %r0 to hold the pid value because there
- * are no assignments to %r0 possible in between the conditional
- * statements.
+ * -sizeof(usdt_prids_map_key_t) + sizeof(pid_t)
+ * ==
+ * -sizeof(dtrace_id_t) underlying-probe prid
*/
- for (pop = dt_list_next(&upp->probes); pop != NULL;
- pop = dt_list_next(pop)) {
- const dt_probe_t *prp = pop->probe;
- pid_t pid;
- dt_ident_t *idp;
-
- pid = dt_pid_get_pid(prp->desc, pcb->pcb_hdl, pcb, NULL);
- assert(pid != -1);
-
- idp = dt_dlib_add_probe_var(pcb->pcb_hdl, prp);
- assert(idp != NULL);
+ emit(dlp, BPF_STORE(BPF_W, BPF_REG_9, (int)(-sizeof(usdt_prids_map_key_t)), BPF_REG_0));
+ emit(dlp, BPF_STORE_IMM(BPF_W, BPF_REG_9, (int)(-sizeof(dtrace_id_t)), uprp->desc->id));
+ dt_cg_xsetx(dlp, usdt_prids, DT_LBL_NONE, BPF_REG_1, usdt_prids->di_id);
+ emit(dlp, BPF_MOV_REG(BPF_REG_2, BPF_REG_9));
+ emit(dlp, BPF_ALU64_IMM(BPF_ADD, BPF_REG_2, (int)(-sizeof(usdt_prids_map_key_t))));
+ emit(dlp, BPF_CALL_HELPER(BPF_FUNC_map_lookup_elem));
+ emit(dlp, BPF_BRANCH_IMM(BPF_JEQ, BPF_REG_0, 0, pcb->pcb_exitlbl));

- /*
- * Check whether this pid-provider probe serves the current
- * process, and copy out a 1 into arg 0 if so.
- */
- emit(dlp, BPF_BRANCH_IMM(BPF_JEQ, BPF_REG_0, pid, lbl_assign));
- }
- emit(dlp, BPF_JUMP(lbl_exit));
- copyout_val(pcb, lbl_assign, 1, 0);
+ /*
+ * If we succeeded, then we use copyout_val() above to assign:
+ * *arg0 = 1;
+ */
+ copyout_val(pcb, DT_LBL_NONE, 1, 0);

dt_cg_tramp_return(pcb);

--
2.43.5