From aad8388a0de70a9efcbdcd583d2dbc28bb49b68e Mon Sep 17 00:00:00 2001 From: 0xHansLee Date: Fri, 17 Jul 2026 12:55:56 +0900 Subject: [PATCH 1/3] fix(dkg): restart decrypt worker whenever an active round exists On a node restart, the TDH2 decrypt worker was only restarted for the latest round; if the next round had already opened registration, the latest round was no longer active and no start path fired, so decrypt requests queued indefinitely. ResumeDKGService now (re)starts the idempotent worker whenever an active DKG round exists, independent of the latest round's stage. --- client/x/dkg/keeper/dkg_svc.go | 8 ++ client/x/dkg/keeper/dkg_svc_test.go | 127 ++++++++++++++++++++++++++++ 2 files changed, 135 insertions(+) diff --git a/client/x/dkg/keeper/dkg_svc.go b/client/x/dkg/keeper/dkg_svc.go index fff42351..ce4d2099 100644 --- a/client/x/dkg/keeper/dkg_svc.go +++ b/client/x/dkg/keeper/dkg_svc.go @@ -100,6 +100,14 @@ var decryptWorkerRunning atomic.Bool // ResumeDKGService reloads unfinished DKG sessions and resumes their execution safely without spawning duplicate goroutines. func (k *Keeper) ResumeDKGService(ctx context.Context, dkgNetwork *types.DKGNetwork) { + // The decrypt worker serves the latest ACTIVE round, which may differ from the + // latest round once the next round has opened. StartDecryptWorker is idempotent. + if active, err := k.getLatestActiveDKGNetwork(ctx); err != nil { + log.Warn(ctx, "Failed to get latest active DKG round while resuming decrypt worker", err) + } else if active != nil { + k.StartDecryptWorker() + } + session, err := k.stateManager.GetSession(dkgNetwork.Round) if err != nil { log.Error(ctx, "Failed to get DKG session while resuming the DKG service", err) diff --git a/client/x/dkg/keeper/dkg_svc_test.go b/client/x/dkg/keeper/dkg_svc_test.go index dcaaa189..e674c5a4 100644 --- a/client/x/dkg/keeper/dkg_svc_test.go +++ b/client/x/dkg/keeper/dkg_svc_test.go @@ -1871,6 +1871,133 @@ func TestResumeDKGService_NoSession(t *testing.T) { k.ResumeDKGService(ctx, dkgNetwork) } +// TestResumeDKGService_ActiveRoundStartsWorker verifies that the decrypt worker is +// restarted whenever an active DKG round exists, even if the latest round has already +// advanced to a pre-active stage (e.g. the next round opened registration). This is +// the restart-window hole from issue piplabs/story#854. +func TestResumeDKGService_ActiveRoundStartsWorker(t *testing.T) { + if testing.Short() { + t.Skip("skipping worker test in short mode") + } + + k, _, _, ctx := setupDKGKeeperWithMocks(t) + + sm, err := NewStateManager(t.TempDir()) + require.NoError(t, err) + k.stateManager = sm + + resetDKGSvcRound() + defer resetDKGSvcRound() + // Force a fresh worker spawn so the assertion exercises the new branch, and + // leave decryptWorkerRunning true afterwards. The worker is a process-lifetime + // singleton; resetting the flag to false here would let a later mock-backed test + // spawn a real worker that calls its gomock client after that test completes. + decryptWorkerRunning.Store(false) + + // Detach the mock contract client so the process-lifetime worker goroutine + // returns early at its nil guard instead of issuing unexpected mock calls + // after the test completes. + k.contractClient = nil + + // An earlier round is still the active (completed) round. + activeNetwork := &types.DKGNetwork{ + Round: 5, + Stage: types.DKGStageActive, + } + require.NoError(t, k.setDKGNetwork(ctx, activeNetwork)) + require.NoError(t, k.setLatestActiveRound(ctx, activeNetwork)) + + // The latest round has opened registration for the NEXT round; its session is in a + // valid (non-stuck) phase, so none of the legacy branches would start the worker. + session := &types.DKGSession{ + Round: 6, + Phase: types.PhaseInitialized, + } + require.NoError(t, sm.CreateSession(ctx, session)) + + dkgNetwork := &types.DKGNetwork{ + Round: 6, + Stage: types.DKGStageRegistration, + } + + k.ResumeDKGService(ctx, dkgNetwork) + + require.True(t, decryptWorkerRunning.Load(), "worker should run whenever an active round exists") +} + +// TestResumeDKGService_NoActiveRoundNoWorker verifies that on a fresh chain with no +// active round, ResumeDKGService does not start the decrypt worker via the new branch. +func TestResumeDKGService_NoActiveRoundNoWorker(t *testing.T) { + k, _, _, ctx := setupDKGKeeperWithMocks(t) + + sm, err := NewStateManager(t.TempDir()) + require.NoError(t, err) + k.stateManager = sm + + resetDKGSvcRound() + defer resetDKGSvcRound() + decryptWorkerRunning.Store(false) + // Restore the singleton flag to true on exit so a later mock-backed test does + // not spawn a real worker that outlives it (see the note above). + defer decryptWorkerRunning.Store(true) + + // Latest round session is in a valid (non-stuck) phase and no active round exists. + session := &types.DKGSession{ + Round: 1, + Phase: types.PhaseInitialized, + } + require.NoError(t, sm.CreateSession(ctx, session)) + + dkgNetwork := &types.DKGNetwork{ + Round: 1, + Stage: types.DKGStageRegistration, + } + + k.ResumeDKGService(ctx, dkgNetwork) + + require.False(t, decryptWorkerRunning.Load(), "worker must not start when no active round exists") +} + +// TestResumeDKGService_CompletedActiveStartsWorker verifies the pre-existing behavior: +// when the latest round is itself active with a completed session, the worker is started. +func TestResumeDKGService_CompletedActiveStartsWorker(t *testing.T) { + if testing.Short() { + t.Skip("skipping worker test in short mode") + } + + k, _, _, ctx := setupDKGKeeperWithMocks(t) + + sm, err := NewStateManager(t.TempDir()) + require.NoError(t, err) + k.stateManager = sm + + resetDKGSvcRound() + defer resetDKGSvcRound() + // Force a fresh spawn for the assertion; leave the flag true on exit (see the + // note in TestResumeDKGService_ActiveRoundStartsWorker). + decryptWorkerRunning.Store(false) + + // Detach the mock contract client so the process-lifetime worker goroutine + // returns early at its nil guard instead of issuing unexpected mock calls + // after the test completes. + k.contractClient = nil + + session := &types.DKGSession{ + Round: 7, + Phase: types.PhaseCompleted, + } + require.NoError(t, sm.CreateSession(ctx, session)) + + dkgNetwork := &types.DKGNetwork{ + Round: 7, + Stage: types.DKGStageActive, + } + + k.ResumeDKGService(ctx, dkgNetwork) + + require.True(t, decryptWorkerRunning.Load(), "worker should run when latest active round session is completed") +} + // --- resumeFailedSession --- func TestResumeFailedSession_DealingStage(t *testing.T) { From b0cba30bf16af0f5b25595c6eaf19a5157966e3e Mon Sep 17 00:00:00 2001 From: 0xHansLee Date: Tue, 21 Jul 2026 18:49:45 +0900 Subject: [PATCH 2/3] feat(dkg): log decrypt worker idle/no-op transitions Add Debug logs at the two spots where the decrypt worker silently does nothing: when there is no active round so the worker is not started, and when the running flag CompareAndSwap no-ops because the worker is already running. Without these, a healthy-but-idle worker cannot be told apart from a stuck flag over a dead goroutine. Debug (not Info) since this runs every block. --- client/x/dkg/keeper/dkg_svc.go | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/client/x/dkg/keeper/dkg_svc.go b/client/x/dkg/keeper/dkg_svc.go index ce4d2099..3395e568 100644 --- a/client/x/dkg/keeper/dkg_svc.go +++ b/client/x/dkg/keeper/dkg_svc.go @@ -106,6 +106,10 @@ func (k *Keeper) ResumeDKGService(ctx context.Context, dkgNetwork *types.DKGNetw log.Warn(ctx, "Failed to get latest active DKG round while resuming decrypt worker", err) } else if active != nil { k.StartDecryptWorker() + } else { + // Debug (not Info): this runs every block; a missing active round is the + // normal steady state and Info would be too noisy. + log.Debug(ctx, "No active DKG round; decrypt worker not started") } session, err := k.stateManager.GetSession(dkgNetwork.Round) @@ -287,7 +291,10 @@ func (k *Keeper) resumeFailedSession(ctx context.Context, session *types.DKGSess // goroutine exits. The decrypt worker must run for the lifetime of the process. func (k *Keeper) StartDecryptWorker() { if !decryptWorkerRunning.CompareAndSwap(false, true) { - // already running + // Debug (not Info): this runs every block; skipping because the worker is + // already running is the normal steady state and Info would be too noisy. + log.Debug(context.Background(), "Decrypt worker already running; skipping start") + return } From 0e512c609ff4c64dcce94f9dca73cbd404b3bdfd Mon Sep 17 00:00:00 2001 From: 0xHansLee Date: Tue, 21 Jul 2026 19:05:33 +0900 Subject: [PATCH 3/3] fix(dkg): bound decrypt worker kernel calls with timeout and keepalive The decrypt worker runs on a process-lifetime context.Background with no deadline, and it was passed straight into the kernel PartialDecryptTDH2 RPC. A wedged kernel (accepts the connection but never responds) blocked the worker's synchronous loop forever, so decryptWorkerRunning was never reset and ResumeDKGService's CompareAndSwap kept failing, leaving the worker dead with no self-heal. Wrap each kernel RPC the worker makes in a context.WithTimeout derived context (decryptKernelCallTimeout, 60s, matching the kernel-op budget in dkgAsyncTimeout and the contract-call timeout), and add gRPC keepalive to the kernel client dial options so a wedged connection is detected at the transport layer. Node-local worker path only; no consensus code touched. --- client/x/dkg/keeper/dkg_svc.go | 15 ++++- client/x/dkg/keeper/dkg_svc_test.go | 99 ++++++++++++++++++++++++++++ client/x/dkg/keeper/kernel_client.go | 17 ++++- 3 files changed, 129 insertions(+), 2 deletions(-) diff --git a/client/x/dkg/keeper/dkg_svc.go b/client/x/dkg/keeper/dkg_svc.go index 3395e568..2b323546 100644 --- a/client/x/dkg/keeper/dkg_svc.go +++ b/client/x/dkg/keeper/dkg_svc.go @@ -90,6 +90,14 @@ func releaseDKGSvc(round uint32) { // is cut short — so this value must stay above 60s to be meaningful. const dkgAsyncTimeout = 2 * time.Minute +// decryptKernelCallTimeout bounds a single kernel RPC made by the decrypt worker. +// The worker runs on a process-lifetime context.Background with no deadline of its +// own, so without a per-call timeout a wedged kernel (accepts the connection but +// never responds) blocks the worker loop forever and decryptWorkerRunning is never +// reset. 60s matches the kernel-operation budget in dkgAsyncTimeout and the +// contract-call timeout in contract_client.go. +const decryptKernelCallTimeout = 60 * time.Second + // dkgAsyncContext creates a new context for async DKG service goroutines with a timeout. // This replaces the consensus context that would otherwise be canceled after block processing. func dkgAsyncContext() (context.Context, context.CancelFunc) { @@ -635,8 +643,13 @@ func (k *Keeper) computePartialDecrypt(ctx context.Context, session *types.DKGSe return } + // Bound the kernel RPC so a wedged kernel cannot block the worker loop + // indefinitely. The worker's context has no deadline of its own. + callCtx, cancel := context.WithTimeout(ctx, decryptKernelCallTimeout) + defer cancel() + kernelStart := time.Now() - resp, err := client.PartialDecryptTDH2(ctx, &types.PartialDecryptTDH2Request{ + resp, err := client.PartialDecryptTDH2(callCtx, &types.PartialDecryptTDH2Request{ CodeCommitment: session.CodeCommitment, Round: session.Round, Ciphertext: req.Ciphertext, diff --git a/client/x/dkg/keeper/dkg_svc_test.go b/client/x/dkg/keeper/dkg_svc_test.go index e674c5a4..30b816f8 100644 --- a/client/x/dkg/keeper/dkg_svc_test.go +++ b/client/x/dkg/keeper/dkg_svc_test.go @@ -2216,3 +2216,102 @@ func TestResumeFailedSession_ActiveStage(t *testing.T) { require.Equal(t, types.PhaseCompleted, got.Phase, "active stage: handleDKGComplete should advance phase to PhaseCompleted") } + +// --- decrypt worker kernel-call deadline (issue piplabs/story#854 family) --- + +// TestComputePartialDecrypt_DerivesCallDeadline verifies that the kernel RPC is +// invoked with a bounded context even when the worker passes context.Background +// (which has no deadline). This proves decryptKernelCallTimeout is wired so a +// wedged kernel cannot block the worker loop indefinitely. +func TestComputePartialDecrypt_DerivesCallDeadline(t *testing.T) { + t.Parallel() + + ctx := context.Background() + ctrl := gomock.NewController(t) + + cc := []byte("cc-deadline") + mockKernel := dkgtestutil.NewMockKernelServiceClient(ctrl) + + router := NewKernelRouter(nil, nil) + router.RegisterClient(cc, mockKernel) + + k := &Keeper{kernelRouter: router} + + session := &types.DKGSession{ + Round: 1, Index: 1, GlobalPubKey: []byte("pub"), CodeCommitment: cc, + } + req := types.DecryptRequest{Ciphertext: []byte("ct"), Label: make([]byte, 32)} + + var ( + gotDeadline bool + remaining time.Duration + ) + mockKernel.EXPECT().PartialDecryptTDH2(gomock.Any(), gomock.Any()). + DoAndReturn(func(callCtx context.Context, _ *types.PartialDecryptTDH2Request, _ ...interface{}) (*types.PartialDecryptTDH2Response, error) { + deadline, ok := callCtx.Deadline() + gotDeadline = ok + if ok { + remaining = time.Until(deadline) + } + return &types.PartialDecryptTDH2Response{ + EncryptedPartialDecryption: []byte("p"), EphemeralPubKey: []byte("e"), + PubShare: []byte("s"), Signature: []byte("sig"), + }, nil + }).Times(1) + + // Parent context has NO deadline; the deadline must come from the fix. + result := k.computePartialDecrypt(ctx, session, req) + require.NoError(t, result.err) + require.True(t, gotDeadline, "kernel call must receive a bounded context even from a deadline-less parent") + require.Greater(t, remaining, time.Duration(0), "derived deadline must be in the future") + require.LessOrEqual(t, remaining, decryptKernelCallTimeout, "derived deadline must not exceed the configured timeout") +} + +// TestComputePartialDecrypt_WedgedKernelReturns simulates a kernel that accepts the +// call but never responds on its own. The worker call must still return (via context +// timeout) rather than block forever, so the worker loop can complete and +// decryptWorkerRunning can be reset. Uses a short-deadline parent context to keep the +// test fast; the production timeout is decryptKernelCallTimeout. +func TestComputePartialDecrypt_WedgedKernelReturns(t *testing.T) { + t.Parallel() + + ctrl := gomock.NewController(t) + + cc := []byte("cc-wedged") + mockKernel := dkgtestutil.NewMockKernelServiceClient(ctrl) + + router := NewKernelRouter(nil, nil) + router.RegisterClient(cc, mockKernel) + + k := &Keeper{kernelRouter: router} + + session := &types.DKGSession{ + Round: 1, Index: 1, GlobalPubKey: []byte("pub"), CodeCommitment: cc, + } + req := types.DecryptRequest{Ciphertext: []byte("ct"), Label: make([]byte, 32)} + + // Kernel never returns on its own: it blocks until the call context is done, + // then surfaces the context error (mirrors a real gRPC deadline cancellation). + mockKernel.EXPECT().PartialDecryptTDH2(gomock.Any(), gomock.Any()). + DoAndReturn(func(callCtx context.Context, _ *types.PartialDecryptTDH2Request, _ ...interface{}) (*types.PartialDecryptTDH2Response, error) { + <-callCtx.Done() + return nil, callCtx.Err() + }).Times(1) + + // Short-deadline parent stands in for decryptKernelCallTimeout to keep the test fast. + parentCtx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + + done := make(chan decryptComputeResult, 1) + go func() { + done <- k.computePartialDecrypt(parentCtx, session, req) + }() + + select { + case result := <-done: + require.Error(t, result.err, "wedged kernel call must return an error, not a response") + require.Contains(t, result.err.Error(), "generating partial decrypt failed") + case <-time.After(5 * time.Second): + t.Fatal("computePartialDecrypt blocked on a wedged kernel; worker would never reset decryptWorkerRunning") + } +} diff --git a/client/x/dkg/keeper/kernel_client.go b/client/x/dkg/keeper/kernel_client.go index b54ae6b1..c9f56efb 100644 --- a/client/x/dkg/keeper/kernel_client.go +++ b/client/x/dkg/keeper/kernel_client.go @@ -6,6 +6,7 @@ import ( "io" "os" "strings" + "time" "github.com/piplabs/story/client/x/dkg/types" "github.com/piplabs/story/lib/errors" @@ -13,8 +14,19 @@ import ( "google.golang.org/grpc" "google.golang.org/grpc/credentials" "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/keepalive" ) +// kernelKeepalive detects a wedged kernel connection (accepts the socket but stops +// responding) at the transport layer so a dead connection surfaces as a call error +// instead of hanging a caller forever. PermitWithoutStream keeps pinging even when +// no RPC is in flight, since the decrypt worker is idle between ticks. +var kernelKeepalive = keepalive.ClientParameters{ + Time: 30 * time.Second, + Timeout: 10 * time.Second, + PermitWithoutStream: true, +} + // TLSConfig holds the TLS certificate paths for gRPC client connections. // When CAFile is set, TLS is enabled with server certificate verification. // When CertFile and KeyFile are also set, mutual TLS (mTLS) is used. @@ -56,7 +68,10 @@ func CreateKernelClient(endpoint string, tlsCfg *TLSConfig) (types.KernelService target = "passthrough:///" + target } - conn, err := grpc.NewClient(target, grpc.WithTransportCredentials(creds)) + conn, err := grpc.NewClient(target, + grpc.WithTransportCredentials(creds), + grpc.WithKeepaliveParams(kernelKeepalive), + ) if err != nil { return nil, nil, errors.Wrap(err, "failed to connect to story-kernel client") }