fix huge memory leak in RPC system

[sawka]

always remove from reqmap in Finalize even if done flag is already set.

[coderabbitai]

Walkthrough

Changes to pkg/wshutil/wshrpc.go introduce a new private helper method handleEventRecv for processing EventRecv data with nil-guards and WaveEvent unmarshalling. The handleRequest method now passes a pprof context through the request handling pipeline. The handleRequestInternal method signature is updated to accept an additional context.Context parameter and routes EventRecv commands through the new helper. The Finalize method adds a conditional guard to unregister response handlers earlier when a request ID is present.

Estimated code review effort

🎯 3 (Moderate) | ⏱️ ~20 minutes

• New private method logic: The handleEventRecv helper introduces event processing logic that requires verification of nil-guards, unmarshalling behavior, and EventListener invocation semantics
• Context propagation: Adding pprof context as a parameter through the request handling chain requires verification that context is threaded correctly through call sites
• Method signature change: The update to handleRequestInternal is a breaking change to a private method that necessitates checking all callers and ensuring the pprof context is correctly populated
• Finalize refactoring: The conditional unregistration logic adds race condition considerations that warrant careful review of the timing and ordering of operations

Pre-merge checks and finishing touches

✅ Passed checks (2 passed)

Check name	Status	Explanation
Title check	✅ Passed	The title directly reflects the main objective of the PR - fixing a memory leak in the RPC system by ensuring proper cleanup.
Description check	✅ Passed	The description is directly related to the changeset, explaining the specific fix: always removing from reqmap in Finalize even if done flag is already set.

✨ Finishing touches

• 📝 Generate docstrings

🧪 Generate unit tests (beta)

• Create PR with unit tests
• Post copyable unit tests in a comment
• Commit unit tests in branch sawka/fix-rpc-memleak

Warning

There were issues while running some tools. Please review the errors and either fix the tool's configuration or disable the tool if it's a critical failure.

🔧 golangci-lint (2.5.0)

Error: unknown linters: 'unusedfunc,unusedparams', run 'golangci-lint help linters' to see the list of supported linters
The command is terminated due to an error: unknown linters: 'unusedfunc,unusedparams', run 'golangci-lint help linters' to see the list of supported linters

---

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

❤️ Share

• X
• Mastodon
• Reddit
• LinkedIn

_{Comment @coderabbitai help to get the list of available commands and usage tips.}

[coderabbitai]

Actionable comments posted: 0

🧹 Nitpick comments (1)

pkg/wshutil/wshrpc.go (1)

277-287: Consider logging unmarshal errors for debugging.

The silent error handling on lines 283-285 might make it difficult to debug malformed event data. While events are fire-and-forget, logging errors would help identify issues without breaking the event flow.

Consider this change:

 func (w *WshRpc) handleEventRecv(req *RpcMessage) {
 	if req.Data == nil {
 		return
 	}
 	var waveEvent wps.WaveEvent
 	err := utilfn.ReUnmarshal(&waveEvent, req.Data)
 	if err != nil {
+		log.Printf("[%s] failed to unmarshal wave event: %v\n", w.DebugName, err)
 		return
 	}
 	w.EventListener.RecvEvent(&waveEvent)
 }

📜 Review details

Configuration used: CodeRabbit UI

Review profile: CHILL

Plan: Pro

📥 Commits

Reviewing files that changed from the base of the PR and between 984ed46 and c833bb6.

⛔ Files ignored due to path filters (1)

• package-lock.json is excluded by !**/package-lock.json

📒 Files selected for processing (1)

• pkg/wshutil/wshrpc.go (2 hunks)

🧰 Additional context used

🧠 Learnings (3)

📓 Common learnings

Learnt from: esimkowitz
Repo: wavetermdev/waveterm PR: 1790
File: pkg/remote/fileshare/wshfs/wshfs.go:122-122
Timestamp: 2025-01-22T01:28:41.417Z
Learning: The RpcClient in pkg/remote/fileshare/wshfs/wshfs.go is initialized and handled downstream by either main-server or wshcmd-connserver, as documented in the package comment.

📚 Learning: 2025-01-22T01:28:41.417Z

Learnt from: esimkowitz
Repo: wavetermdev/waveterm PR: 1790
File: pkg/remote/fileshare/wshfs/wshfs.go:122-122
Timestamp: 2025-01-22T01:28:41.417Z
Learning: The RpcClient in pkg/remote/fileshare/wshfs/wshfs.go is initialized and handled downstream by either main-server or wshcmd-connserver, as documented in the package comment.

Applied to files:

• pkg/wshutil/wshrpc.go

📚 Learning: 2025-01-22T22:27:25.739Z

Learnt from: esimkowitz
Repo: wavetermdev/waveterm PR: 1790
File: pkg/remote/connparse/connparse.go:76-82
Timestamp: 2025-01-22T22:27:25.739Z
Learning: The GetRpcContext() method in wshutil package is guaranteed to never return nil due to type constraints, making nil checks unnecessary.

Applied to files:

• pkg/wshutil/wshrpc.go

🧬 Code graph analysis (1)

pkg/wshutil/wshrpc.go (4)

pkg/wps/wpstypes.go (1)

• WaveEvent (25-31)

pkg/util/utilfn/marshal.go (1)

• ReUnmarshal (36-42)

pkg/wshutil/wshevent.go (1)

• EventListener (20-23)

pkg/wshrpc/wshrpctypes.go (1)

• Command_EventRecv (88-88)

⏰ Context from checks skipped due to timeout of 90000ms. You can increase the timeout in your CodeRabbit configuration to a maximum of 15 minutes (900000ms). (3)

• GitHub Check: Build for TestDriver.ai
• GitHub Check: Analyze (go)
• GitHub Check: Analyze (javascript-typescript)

🔇 Additional comments (3)

pkg/wshutil/wshrpc.go (3)

272-274: Good addition for observability.

The pprof labeling by RPC command will enable better performance profiling and debugging of RPC handlers.

---

289-291: Clean refactor of event handling.

Extracting EventRecv logic to a dedicated helper improves code organization and readability.

---

674-677: Critical fix: Resolves the memory leak.

This change correctly ensures that response handlers are always removed from ResponseHandlerMap, even when the done flag is already set. Previously, the early return on line 678 (when done.Load() was true) would skip unregistering, causing handlers to accumulate indefinitely in the map.

The fix is thread-safe: unregisterResponseHandler uses a lock, and multiple calls are idempotent (delete from map is safe even if key doesn't exist).