fix: implemented graceful shutdown for workloadmanager

[mahil-2040]

What type of PR is this?

/kind enhancement

What this PR does / why we need it:
Implements a comprehensive graceful shutdown mechanism for the workloadmanager.

Previously, process termination was abrupt — persistent connections to the backend store (Redis/Valkey) remained open until timeout, and background workers could be interrupted mid-operation.

This PR introduces a coordinated shutdown sequence:

• HTTP server stops accepting new requests and drains in-flight ones via Server.Shutdown(ctx)
• Store connections (Redis/Valkey) are explicitly closed via a new Close() method on the Store interface
• Background workers (GC, informers, controller manager) stop gracefully when the root context is cancelled

---

Changes

• Added Close() error to the Store interface, implemented by both:
  • redisStore
  • valkeyStore
• Added Shutdown(ctx) method to the Server struct for coordinated shutdown
• Updated main.go to call server.Shutdown() with a 15s deadline instead of time.Sleep(2s)
• Fixed fakeStoreClient in session_manager_test.go to satisfy the updated interface

Which issue(s) this PR fixes:
Fixes #215

Special notes for your reviewer:

The shutdown order is intentional:

1. HTTP server drains first
2. Store connections are closed
3. Root context is cancelled to stop background workers (GC, informers)

This prevents requests from failing due to a prematurely closed store.

Does this PR introduce a user-facing change?:

None

[volcano-sh-bot]

[APPROVALNOTIFIER] This PR is NOT APPROVED

This pull-request has been approved by:
Once this PR has been reviewed and has the lgtm label, please assign yaozengzeng for approval. For more information see the Kubernetes Code Review Process.

The full list of commands accepted by this bot can be found here.

Details

Needs approval from an approver in each of these files:

• cmd/OWNERS
• pkg/OWNERS

Approvers can indicate their approval by writing /approve in a comment
Approvers can cancel approval by writing /approve cancel in a comment

[gemini-code-assist]

Summary of Changes

Hello @mahil-2040, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the workloadmanager's resilience and resource management by introducing a robust graceful shutdown process. It ensures that upon termination, the HTTP server properly drains in-flight requests, persistent connections to backend stores (Redis/Valkey) are explicitly closed, and background workers are stopped in a coordinated manner. This prevents data loss, resource leaks, and abrupt interruptions, leading to a more stable and predictable application lifecycle.

Highlights

• Graceful Shutdown Implementation: Implemented a comprehensive graceful shutdown mechanism for the workloadmanager to ensure proper resource release and prevent abrupt terminations.
• Store Interface Enhancement: Introduced a Close() method to the Store interface, which is now implemented by both redisStore and valkeyStore to explicitly release backend store connections.
• Server Shutdown Orchestration: Added a Shutdown(ctx) method to the Server struct, responsible for orchestrating the HTTP server draining and subsequent store connection closure.
• Main Application Logic Update: Updated the main application entry point (main.go) to utilize the new Server.Shutdown() with a 15-second deadline, replacing a fixed time.Sleep.
• Test Compatibility: Ensured test compatibility by adding a Close() method to the fakeStoreClient in session_manager_test.go to satisfy the updated Store interface.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

• cmd/workload-manager/main.go
  • Replaced a fixed time.Sleep with a structured shutdown sequence that uses a 15-second context deadline for server.Shutdown() and then cancels the root context for background workers.
• pkg/router/session_manager_test.go
  • Added a Close() method to the fakeStoreClient to align with the updated Store interface.
• pkg/store/interface.go
  • Defined a new Close() error method in the Store interface for resource cleanup.
• pkg/store/store_redis.go
  • Implemented the Close() method for redisStore to close the underlying Redis client.
• pkg/store/store_valkey.go
  • Implemented the Close() method for valkeyStore to close the underlying Valkey client.
• pkg/workloadmanager/server.go
  • Removed the previous in-goroutine shutdown logic.
  • Introduced a new Shutdown(ctx) method that gracefully shuts down the HTTP server and then closes the store connections.

Activity

• The pull request was created by mahil-2040 to address issue implement a comprehensive graceful shutdown mechanism for the workloadmanager #215.
• The author provided a detailed description outlining the problem, the solution, and specific changes, including a note for reviewers about the intentional shutdown order.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces a graceful shutdown mechanism for the workloadmanager, which is a great improvement over the previous abrupt termination. The changes include adding a Close() method to the Store interface and its implementations, and a new Shutdown() method on the Server. My main feedback concerns a potential race condition in the shutdown order, where the store connection might be closed before background workers using it have terminated. I've left a detailed comment with a suggestion for a safer shutdown sequence.

[Copilot]

Pull request overview

This PR implements a comprehensive graceful shutdown mechanism for the workloadmanager component. Previously, process termination was abrupt—HTTP server shutdown was handled inline, but Redis/Valkey connections remained open until timeout, and background workers could be interrupted mid-operation. This change introduces a coordinated shutdown sequence where the HTTP server drains first, then store connections are closed, and finally background workers are signaled to stop via context cancellation.

Changes:

• Added Close() error method to the Store interface and implemented it in both redisStore and valkeyStore
• Replaced the inline shutdown goroutine in server.Start() with a dedicated Shutdown() method on the Server struct
• Updated main.go to call server.Shutdown() with a 15-second deadline instead of a 2-second sleep

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
pkg/store/interface.go	Added Close() method to Store interface for releasing resources
pkg/store/store_redis.go	Implemented Close() that calls redis client's Close() method
pkg/store/store_valkey.go	Implemented Close() that calls valkey client's Close() method
pkg/workloadmanager/server.go	Removed inline shutdown goroutine and added dedicated Shutdown() method
cmd/workload-manager/main.go	Updated signal handling to use server.Shutdown() with 15s timeout
pkg/router/session_manager_test.go	Updated fakeStoreClient to implement Close() method

[codecov-commenter]

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 0% with 29 lines in your changes missing coverage. Please review.
✅ Project coverage is 41.70%. Comparing base (845b798) to head (44898eb).
⚠️ Report is 99 commits behind head on main.

Files with missing lines	Patch %	Lines
pkg/workloadmanager/server.go	0.00%	24 Missing ⚠️
pkg/store/store_valkey.go	0.00%	3 Missing ⚠️
pkg/store/store_redis.go	0.00%	2 Missing ⚠️
❗ Your organization needs to install the Codecov GitHub app to enable full functionality.

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #217      +/-   ##
==========================================
+ Coverage   35.60%   41.70%   +6.09%     
==========================================
  Files          29       30       +1     
  Lines        2533     2611      +78     
==========================================
+ Hits          902     1089     +187     
+ Misses       1505     1390     -115     
- Partials      126      132       +6     

Flag	Coverage Δ
unittests	41.70% <0.00%> (+6.09%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[mahil-2040]

@hzxuzhonghu PTAL!

[Copilot]

Pull request overview

Copilot reviewed 6 out of 6 changed files in this pull request and generated 10 comments.

[mahil-2040]

@hzxuzhonghu please check this!