gpu: fire the sampled probe before the batched add

gpuprobe/batch_size_test.go

@@ -4,6 +4,7 @@ import (
 	"os"
 	"regexp"
 	"strconv"
+	"strings"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -124,6 +125,43 @@ func TestSampledProbesAreNotBatchedInTheStub(t *testing.T) {
 	}
 }
 
+// Issue #67, for the producer no unprivileged test can execute.
+//
+// The sampled record and the batched gpu_launch_v1 record for one launch are
+// twins, and the consumer's join is only safe when the sampled one leads:
+// batched-first leaves the launch in deferredLaunches, where the next batch
+// of any other kind releases it stackless, and the stack then parks with
+// nothing to join (sampledstacks.go). Both shims therefore fire the sampled
+// probe BEFORE the batched add().
+//
+// For shim/stub/stub.cc that order is proven rather than asserted -
+// shim/stub/probe_order_test.cc patches the probe sites with int3 and reads
+// the wire order back, and it fails on the pre-#67 producer. The CUPTI
+// adapter cannot be driven without a CUDA process and a GPU, so the same
+// fact is pinned against its source here. A regex over C++ is a weak
+// instrument; it is here because the alternative for this file is nothing at
+// all, and because a reorder is exactly the kind of edit that looks harmless
+// in review.
+func TestBothShimsFireTheSampledProbeBeforeTheBatchedAdd(t *testing.T) {
+	for _, tc := range []struct{ file, add string }{
+		{"../shim/stub/stub.cc", "lb.add(l);"},
+		{"../shim/nvidia/cupti_adapter.cc", "g_lb->add(l);"},
+	} {
+		b, err := os.ReadFile(tc.file)
+		require.NoError(t, err)
+		src := string(b)
+		emit := strings.Index(src, "gpu_launch_sampled_v1_emit(&")
+		batched := strings.Index(src, tc.add)
+		require.Positivef(t, emit, "%s: no unbatched sampled-launch emit found", tc.file)
+		require.Positivef(t, batched, "%s: no %q found", tc.file, tc.add)
+		require.Lessf(t, emit, batched,
+			"%s: the launch is added to its batch before the sampled probe fires. "+
+				"A launch that fills the batch then reaches the consumer ahead of its own "+
+				"stack, is released stackless by the next exec batch, and the stack parks "+
+				"in PendingStacks forever (issue #67)", tc.file)
+	}
+}
+
 // The sampled-launch cap is 1, and not because of bytes: the batch header
 // holds one stack id, so a batch of N sampled launches would attribute one
 // captured stack to N unrelated launches. Its byte budget would allow 54.

gpuprobe/consumer_test.go

@@ -1616,11 +1616,11 @@ func TestSampledStackArrivingFirstAttachesToTheBatchedLaunch(t *testing.T) {
 	assert.True(t, sm.wasDeleted(5), "a consumed stackmap entry must be freed")
 }
 
-// The other arrival order, which happens whenever the launch that fills the
-// shim's batch is also the one the sampler picked: the flush is queued
-// inside the add() that precedes the sampler check, so the batched record
-// reaches the ringbuf first. The launch waits for the twin already on its
-// way, then goes out once - with its stack.
+// The other arrival order. No shim in this repository produces it any more -
+// they fire the sampled probe before the batched add(), see issue #67 and
+// sampledstacks.go - but the ABI is public and a foreign or older producer
+// may still emit batched-first, so the consumer must join it. The launch
+// waits for the twin, then goes out once, with its stack.
 func TestBatchedLaunchArrivingFirstWaitsForItsStack(t *testing.T) {
 	sink := &recordingSink{}
 	c, sm, _ := stackConsumer(t, sink, Config{})
@@ -1753,36 +1753,35 @@ func TestHeldLaunchDoesNotTakeAnotherProcessesStack(t *testing.T) {
 		"the pid-4242 launch was released before its own stack arrived: no stack is correct, a borrowed one is not")
 }
 
-// The measured cause of the unattached stacks in cmd/gpu-stub-profile's
-// 2000-launch run (PendingStacks:24 of 250 captured), reproduced here with
-// no privileges and no attach - just the batch order the stub actually
-// produces.
+// What a batched-first producer costs, and the reason no shim here is one.
 //
-// The stub adds to the launch batch, then to the exec batch, then fires the
-// unbatched sampled probe. When one launch is both the record that FILLS
-// the launch batch and the one the sampler picks, all three land on the
-// ringbuf in that order: launch batch, exec batch, sampled record. The
-// launch is held for its twin; the exec batch releases it stackless (which
-// is correct - the timeline needs launches promptly); and the twin then
-// arrives with nowhere to go and parks forever.
+// Three records land in this order: launch batch, exec batch, sampled
+// record. The launch is held for its twin; the exec batch releases it
+// stackless (which is correct - the timeline wants launches promptly); the
+// twin then arrives with nowhere to go and parks forever.
 //
-// With 32-record batches and a period of 8 those two conditions are
-// disjoint by arithmetic - the batch fills at launch i = 0 mod 32 and the
-// sampler picks i = 1 mod 8 - which is why a short run never shows it. What
-// breaks the arithmetic is the stub's Drainer: it flushes both partial
-// batches every 100ms, which RE-PHASES the batch boundary to wherever the
-// loop happened to be. Roughly one tick in eight leaves the new boundary on
-// a sampled launch, and from then until the next tick every one of the
-// remaining fills - one per 32 launches - loses its stack this way. That is
-// the whole rate dependence: a 500-launch run has about one tick and
-// usually shows none, a 2000-launch run has several and showed 24.
+// This WAS the shipped stub. It added to the launch batch, then to the exec
+// batch, then fired the unbatched sampled probe, so a launch that both
+// FILLED the batch and was sampled produced exactly this sequence: 58
+// sampled, 57 attached, 1 parked, on the privileged gate (issue #67). Issue
+// #50's jittered stride is what made a sampled ordinal able to land on a
+// batch boundary at all - at the old fixed stride of 8 the collision was
+// arithmetically impossible against 32-record batches.
 //
-// It costs attribution only. The launch ships, the execution ships, the GPU
-// time is measured and projects as unattributed, and the stack is counted
-// in PendingStacks rather than vanishing. Fixing it would mean holding
-// launches past the next batch, which trades a rare attribution gain for a
-// systematic delay in launch delivery that the timeline's join depends on -
-// a worse trade, so this is documented and counted, not "fixed".
+// The fix is in the producer: both shims now fire the sampled probe BEFORE
+// the batched add(), which makes sampled-first unconditional, and
+// shim/stub/probe_order_test.cc pins that order without any privilege. This
+// test therefore no longer describes our producers. It stays because the
+// ABI is public and a foreign or older producer may still emit in this
+// order, and because what the consumer does then must be a documented,
+// counted outcome rather than a surprise: the launch ships, the execution
+// ships, the GPU time is measured and projects as unattributed, and the
+// stack is counted in PendingStacks rather than vanishing.
+//
+// It is not fixable on this side without giving something up. Holding the
+// launch past the exec batch delays every launch systematically to buy back
+// a rare attribution, and attaching the stack after the fact means
+// re-emitting a launch the sink has already been given.
 func TestStackParksUnattachedWhenAnotherBatchSplitsTheTwins(t *testing.T) {
 	sink := &recordingSink{}
 	c, sm, _ := stackConsumer(t, sink, Config{})
@@ -1818,6 +1817,62 @@ func TestStackParksUnattachedWhenAnotherBatchSplitsTheTwins(t *testing.T) {
 	assert.Zero(t, st.StacksEvicted, "nothing was pushed out; it is still parked")
 }
 
+// The same batch boundary, in the order the fixed producers actually emit
+// (issue #67): the sampled record for the launch that fills the batch goes
+// out BEFORE the batched add(), so it reaches the ringbuf ahead of the batch
+// that carries its twin.
+//
+// Correlation 7 is the record that fills the launch batch here, and the exec
+// batch of that same producer loop iteration follows immediately - the batch
+// that released the launch stackless in the test above. The stack is parked
+// rather than held, so it is not the deferred queue's to release, and the
+// join survives the exec batch untouched. That is the whole difference the
+// reorder buys, stated in the consumer's own vocabulary rather than the
+// producer's; shim/stub/probe_order_test.cc is what pins the producer to
+// this order.
+func TestStackSurvivesASplittingBatchWhenTheSampledRecordLeads(t *testing.T) {
+	sink := &recordingSink{}
+	c, sm, _ := stackConsumer(t, sink, Config{})
+	sm.put(1, 0x1000)
+
+	// Sampled first: the producer fires this probe before the add() that
+	// flushes the batch below.
+	apply(t, c, sampledBatchWith(4242, 7, 1, 8))
+	require.Equal(t, 1, c.Stats().PendingStacks, "the stack waits for its twin")
+
+	// The batch the twin fills, then the exec batch of the same iteration.
+	apply(t, c, launchBatchWith(4242, 5, 6, 7))
+	execs := make([]byte, batchHdrSize+gpuabi.SizeExec)
+	putU32(execs[0:], kindExec)
+	putU32(execs[4:], 1)
+	putU32(execs[16:], 4242)
+	putU64(execs[24:], gpuabi.SizeExec)
+	putU64(execs[batchHdrSize:], 7)
+	apply(t, c, execs)
+	c.Flush()
+
+	st := c.Stats()
+	assert.Equal(t, uint64(1), st.StacksResolved)
+	assert.Equal(t, uint64(1), st.StacksAttached,
+		"the launch that filled the batch must still carry its stack")
+	assert.Zero(t, st.PendingStacks,
+		"nothing may be left parked: the exec batch cannot take a stack, only a held launch")
+	assert.Zero(t, st.StacksEvicted)
+	require.Len(t, sink.launches, 3)
+	// By correlation, not by position: a launch that collects a parked stack
+	// is emitted on the spot while its batch-mates are still held, so
+	// correlation 7 overtakes 5 and 6 here. admitLaunchLocked documents that
+	// reordering and bounds it to one batch.
+	byCorr := map[string][]string{}
+	for _, l := range sink.launches {
+		byCorr[l.Correlation.Value] = frameNames(l.Launch.CPUStack)
+	}
+	assert.Equal(t, []string{"fn_1000"}, byCorr["7"],
+		"correlation 7 is the sampled one and must carry its own stack")
+	assert.Empty(t, byCorr["5"], "an unsampled batch-mate must stay stackless")
+	assert.Empty(t, byCorr["6"], "an unsampled batch-mate must stay stackless")
+}
+
 // A held launch is waiting for a record that would be the next ringbuf
 // sample. Anything else arriving ends the wait: launches must not sit
 // behind an exec batch, because the timeline joins executions against

gpuprobe/gate_test.go

@@ -331,6 +331,35 @@ func TestStubDrivesThePipelineToPprofWithoutAGPU(t *testing.T) {
 	assert.Zero(t, stats.StackWalkScratchFailed, "a per-CPU scratch lookup at key 0 cannot fail on a loaded program")
 	assert.Zero(t, stats.StacksEvicted,
 		"the parked-stack side table must never overflow at this launch rate and capacity")
+	// Issue #67, and the assertion that keeps it from regressing silently.
+	// Every sampled record now reaches the ringbuf before the batched twin it
+	// belongs to (shim/stub/stub.cc fires the probe before the batched add(),
+	// pinned unprivileged by shim/stub/probe_order_test.cc), so every parked
+	// stack has a launch still to come and the run is flushed by the time
+	// this reads. A stack left parked at rest is a stack whose twin was
+	// emitted before it and released stackless: attribution lost, silently
+	// except for this gauge. It read 1 of 58 on main at 75ecc513.
+	//
+	// Checked before StacksAttached below because it is the more specific
+	// failure: a shortfall in attached stacks could come from a dozen places,
+	// a non-zero PendingStacks names one.
+	assert.Zero(t, stats.PendingStacks,
+		"a resolved stack is still parked with no launch to join, at rest and after Flush: sampled=%d resolved=%d attached=%d evicted=%d profiler-only=%d uncorrelated=%d",
+		stats.SampledLaunches, stats.StacksResolved, stats.StacksAttached,
+		stats.StacksEvicted, stats.StacksProfilerOnly, stats.StacksUncorrelated)
+	assert.Equal(t, uint64(wantSampled), stats.StacksAttached,
+		"every sampled launch must reach the timeline carrying its own stack; resolved=%d pending=%d evicted=%d profiler-only=%d uncorrelated=%d missing=%d",
+		stats.StacksResolved, stats.PendingStacks, stats.StacksEvicted,
+		stats.StacksProfilerOnly, stats.StacksUncorrelated, stats.StacksMissing)
+	// The whole join, on one line, in the shape the accounting identity above
+	// StacksResolved is written in: resolved = attached + evicted +
+	// profiler-only + pending. Printed whether or not the assertions pass,
+	// because "58 = 57 + 0 + 0 + 1" is what told issue #67 apart from silent
+	// loss in the first place.
+	t.Logf("stack attach: sampled=%d resolved=%d attached=%d evicted=%d profiler-only=%d pending=%d missing=%d uncorrelated=%d",
+		stats.SampledLaunches, stats.StacksResolved, stats.StacksAttached,
+		stats.StacksEvicted, stats.StacksProfilerOnly, stats.PendingStacks,
+		stats.StacksMissing, stats.StacksUncorrelated)
 	assert.Zero(t, stats.StackLookupFailed,
 		"every resolved stack's gpu_stacks entry must be readable back exactly once")
 	assert.Zero(t, stats.StackDeleteFailed, "every gpu_stacks entry read must also be deletable")

gpuprobe/sampledstacks.go

@@ -16,46 +16,54 @@ import (
 // launch that arrives on the *batched* probe, and only that launch is
 // emitted.
 //
-// The two halves can arrive in either order, and both happen:
+// The two halves can arrive in either order, and the consumer handles both:
 //
-//   - Sampled first (the common case). The shim's launch batch flushes only
-//     when it fills, so the batched record for launch N usually reaches the
-//     ringbuf long after the unbatched sampled record for the same launch.
-//     The resolved stack waits in pendingStacks until its twin shows up.
-//   - Batched first. When launch N is the record that fills the batch, the
-//     flush happens inside the same add() call that precedes the sampler
-//     check, so the batch - whose last record is launch N - is queued before
-//     N's own sampled record. The launch waits in deferredLaunches, briefly,
-//     for the twin that is already on its way.
+//   - Sampled first. The resolved stack waits in pendingStacks until its
+//     twin shows up. Nothing can take it but the twin, and any number of
+//     unrelated batches may pass in the meantime.
+//   - Batched first. The launch waits in deferredLaunches for a twin that
+//     had better be the very next thing off the ringbuf, because the first
+//     batch of any other kind releases the whole queue (Consumer.applyBatch,
+//     deliberately: the timeline wants launches promptly).
 //
-// The batched-first case has a known, counted, rate-dependent way of losing
-// the join, and it is the cause of the PendingStacks the demo run reports
-// (24 of 250 captures at 2000 launches; none at 500). The producer queues
-// the launch batch, then the EXEC batch, then the sampled probe. So when one
-// launch is both the record that fills the launch batch and the one the
-// sampler picks, the exec batch lands between the twins - and any batch that
-// is not a sampled launch releases the deferred queue (Consumer.applyBatch,
-// deliberately: the timeline needs launches promptly). The launch goes out
-// stackless and its stack, arriving next, parks with nothing to join.
+// Only the first order is safe, and issue #67 is the measurement of that.
+// The shim used to add the launch to its batch and only then fire the
+// sampled probe, so a launch that both FILLED the batch and was sampled put
+// its own batched record on the wire first - with the exec batch of the same
+// loop iteration landing between the twins. The exec batch released the
+// launch stackless and the stack, arriving next, parked with nothing to
+// join: 58 sampled, 57 attached, 1 in PendingStacks on the privileged gate.
 //
-// It looks arithmetically impossible in the stub, which is why it went
-// unexplained: batches hold 32 records and the sampler takes one in 8, so
-// "fills the batch" (i = 0 mod 32) and "is sampled" (i = 1 mod 8) cannot
-// both hold. What breaks that is the producer's periodic drain tick, which
-// flushes a PARTIAL batch every 100ms and so re-phases the batch boundary to
-// wherever the launch loop had got to. About one tick in eight leaves the
-// new boundary sitting on a sampled launch, and until the next tick moves it
-// again every remaining fill - one per 32 launches - loses its stack this
-// way. Hence rate-dependent: more launches, more ticks, more chances to land
-// in phase and more fills to spend there.
+// At the old fixed sampler stride that collision was arithmetically
+// unreachable - batches hold 32 records and the sampler took one in 8, and a
+// multiple of 8 is never 31 mod 32 - which is why the join looked sound for
+// a whole phase. Issue #50's jittered stride draws each gap from [4,12], so
+// a sampled ordinal eventually lands exactly on a batch boundary. #50 did
+// not cause this; it removed the arithmetic that was hiding it.
 //
-// The cost is attribution, never a record: the launch ships, the execution
+// The fix is in the producer, not here: shim/stub/stub.cc and
+// shim/nvidia/cupti_adapter.cc fire the sampled probe BEFORE the batched
+// add(). A record cannot be in a batch before add() puts it there, so no
+// flush - on the launching thread or on the drain thread - can carry the
+// twin past the sampled probe, and sampled-first holds unconditionally.
+// shim/stub/probe_order_test.cc pins that order by patching the probe sites
+// with int3 and reading the wire order back, with no privilege and no
+// consumer; it fails on the pre-#67 producer at every sampled launch that
+// fills a batch.
+//
+// deferredLaunches stays, because this consumer does not only see shims this
+// repository builds: the ABI is public (spec §6), a vendor bridge or an
+// older shim may still emit batched-first, and for those the queue is the
+// difference between "usually joins" and "never joins". What it cannot do is
+// make batched-first lossless - the launch has to go out before the next
+// exec batch, and once it has gone out there is nothing left to attach to
+// without re-emitting an event the sink has already been given. So the cost
+// there is attribution and never a record: the launch ships, the execution
 // ships, the GPU time is measured and projects as unattributed, and the
-// orphaned stack is visible in Stats.PendingStacks. The alternative -
-// holding launches past the next batch - trades a rare attribution gain for
-// a systematic delay in launch delivery that the timeline's join depends on.
-// Reproduced without any privilege by
-// TestStackParksUnattachedWhenAnotherBatchSplitsTheTwins.
+// orphaned stack is counted in Stats.PendingStacks rather than vanishing.
+// That path is held to its documented behaviour by
+// TestStackParksUnattachedWhenAnotherBatchSplitsTheTwins, which is now a
+// statement about foreign producers rather than about ours.
 //
 // Both stores are bounded, and both count what they push out. An unbounded
 // map on either side is a leak driven by a profiled application's launch