Refactor lookup client/server and abstract rpc layer. by maobaolong · Pull Request #2609 · LMCache/LMCache

maobaolong · 2026-02-17T05:36:13Z

What this PR does / why we need it:

This pull request significantly refactors the RPC communication within the LMCache lookup client and server. By introducing an abstract transport layer, the core lookup logic is now decoupled from the specific ZMQ implementation, enhancing modularity and maintainability. This change centralizes ZMQ-specific concerns like socket management, serialization, and error handling into dedicated transport classes, making the client and server components cleaner and more focused on their primary responsibilities.

Special notes for your reviewers:

If applicable:

this PR contains user facing changes - docs added
this PR contains unit tests

gemini-code-assist · 2026-02-17T05:36:34Z

Summary of Changes

Hello @maobaolong, 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 refactors the RPC communication within the LMCache lookup client and server. By introducing an abstract transport layer, the core lookup logic is now decoupled from the specific ZMQ implementation, enhancing modularity and maintainability. This change centralizes ZMQ-specific concerns like socket management, serialization, and error handling into dedicated transport classes, making the client and server components cleaner and more focused on their primary responsibilities.

Highlights

RPC Layer Abstraction: The RPC communication layer for the lookup client and server has been abstracted, decoupling the core logic from the underlying communication mechanism.
New RPC Transport Modules: Introduced new modules lmcache/v1/rpc/transport.py for abstract RPC interfaces and lmcache/v1/rpc/zmq_transport.py for ZMQ-specific implementations.
Dependency Injection for Transports: The LMCacheLookupClient and LMCacheLookupServer now accept an RpcClientTransport or RpcServerTransport object in their constructors, respectively.
Simplified Client/Server Logic: Direct ZMQ socket handling, timeout management, and socket recreation logic have been removed from LMCacheLookupClient and LMCacheLookupServer, delegating these responsibilities to the new transport layer.
Factory Integration: The LookupClientFactory has been updated to create and inject the appropriate ZMQ transport instances when instantiating lookup clients and servers.

🧠 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

lmcache/v1/lookup_client/factory.py
- Imported new ZMQ transport classes and RPC utility functions.
- Modified create_lookup_client to instantiate LMCacheLookupClient with a transport object.
- Modified create_lookup_server to instantiate LMCacheLookupServer with a transport object.
- Added static methods _create_zmq_client_transport and _create_zmq_server_transport for ZMQ transport creation.
lmcache/v1/lookup_client/lmcache_lookup_client.py
- Removed direct ZMQ imports and related socket management logic.
- Updated LMCacheLookupClient and LMCacheLookupServer constructors to accept abstract RpcClientTransport and RpcServerTransport instances.
- Delegated all send_multipart, recv, and close operations to the injected transport objects.
- Moved msgspec.msgpack.Encoder and Decoder usage to the transport layer.
- Adjusted lookup method to interact with the transport for sending and receiving data.
lmcache/v1/rpc/init.py
- Added new file to export RpcClientTransport, RpcServerTransport, ZmqReqRepClientTransport, and ZmqRouterServerTransport.
lmcache/v1/rpc/transport.py
- Added new file defining abstract base classes RpcClientTransport and RpcServerTransport.
- Defined abstract methods for sending/receiving requests and managing resources for RPC communication.
lmcache/v1/rpc/zmq_transport.py
- Added new file implementing ZmqReqRepClientTransport and ZmqRouterServerTransport using ZMQ sockets.
- Encapsulated ZMQ socket creation, configuration, send/recv logic, timeout handling, and automatic socket recreation.
- Implemented msgspec.msgpack serialization/deserialization for message frames within the transport classes.
tests/v1/lookup_client/test_lmcache_lookup_client.py
- Removed direct ZMQ import.
- Imported LookupClientFactory.
- Added helper methods _create_server and _create_client to encapsulate the creation of lookup client/server with their respective transports using the factory.
- Updated all test cases to use these new helper methods for instantiating LMCacheLookupClient and LMCacheLookupServer.
- Removed assertions related to direct ZMQ socket properties.

Activity

The pull request description does not contain specific activity details or comments from reviewers.

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.

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 significantly refactors the RPC layer for the lookup client and server, improving maintainability and extensibility. However, a critical Denial of Service vulnerability was identified in the RPC server's request processing loop due to a lack of error handling for malformed messages, which can crash the background server thread. Further enhancements are suggested to ensure consistent serialization/deserialization in the transport layer and to fix a potential bug in request validation.

maobaolong · 2026-02-19T11:37:16Z

@sammshen @chunxiaozheng Would you like to take a look at this PR? Thanks!

sammshen

LGTM

chunxiaozheng

lgtm

Signed-off-by: baoloongmao <baoloongmao@tencent.com>

Address gemini-bot review comments: 1. Fix critical DoS vulnerability in RPC server request processing - Add comprehensive error handling for malformed messages - Handle JSONDecodeError, UnicodeDecodeError, and other exceptions - Add frame structure and data type validation - Support both bytes and str for backward compatibility 2. Fix request validation bug in transport layer - Change minimum frame count from 2 to 3 - Prevent edge case when data_frames[0] == data_frames[-2] 3. Improve serialization abstraction - Remove redundant UTF-8 encoding in client - Let transport layer handle all serialization - Pass Python strings directly instead of bytes Signed-off-by: maobaolong <maobaolong@tencent.com> Signed-off-by: baoloongmao <baoloongmao@tencent.com>

* Refactor lookup client/server and abstract rpc layer. Signed-off-by: baoloongmao <baoloongmao@tencent.com> * Fix security vulnerabilities and improve RPC robustness Address gemini-bot review comments: 1. Fix critical DoS vulnerability in RPC server request processing - Add comprehensive error handling for malformed messages - Handle JSONDecodeError, UnicodeDecodeError, and other exceptions - Add frame structure and data type validation - Support both bytes and str for backward compatibility 2. Fix request validation bug in transport layer - Change minimum frame count from 2 to 3 - Prevent edge case when data_frames[0] == data_frames[-2] 3. Improve serialization abstraction - Remove redundant UTF-8 encoding in client - Let transport layer handle all serialization - Pass Python strings directly instead of bytes Signed-off-by: maobaolong <maobaolong@tencent.com> Signed-off-by: baoloongmao <baoloongmao@tencent.com> --------- Signed-off-by: baoloongmao <baoloongmao@tencent.com> Signed-off-by: maobaolong <maobaolong@tencent.com> Signed-off-by: shaoxiawjc <wjc2800@163.com>

* Refactor lookup client/server and abstract rpc layer. Signed-off-by: baoloongmao <baoloongmao@tencent.com> * Fix security vulnerabilities and improve RPC robustness Address gemini-bot review comments: 1. Fix critical DoS vulnerability in RPC server request processing - Add comprehensive error handling for malformed messages - Handle JSONDecodeError, UnicodeDecodeError, and other exceptions - Add frame structure and data type validation - Support both bytes and str for backward compatibility 2. Fix request validation bug in transport layer - Change minimum frame count from 2 to 3 - Prevent edge case when data_frames[0] == data_frames[-2] 3. Improve serialization abstraction - Remove redundant UTF-8 encoding in client - Let transport layer handle all serialization - Pass Python strings directly instead of bytes Signed-off-by: maobaolong <maobaolong@tencent.com> Signed-off-by: baoloongmao <baoloongmao@tencent.com> --------- Signed-off-by: baoloongmao <baoloongmao@tencent.com> Signed-off-by: maobaolong <maobaolong@tencent.com> Signed-off-by: Aaron Wu <aaron.wu@dell.com>

* Refactor lookup client/server and abstract rpc layer. Signed-off-by: baoloongmao <baoloongmao@tencent.com> * Fix security vulnerabilities and improve RPC robustness Address gemini-bot review comments: 1. Fix critical DoS vulnerability in RPC server request processing - Add comprehensive error handling for malformed messages - Handle JSONDecodeError, UnicodeDecodeError, and other exceptions - Add frame structure and data type validation - Support both bytes and str for backward compatibility 2. Fix request validation bug in transport layer - Change minimum frame count from 2 to 3 - Prevent edge case when data_frames[0] == data_frames[-2] 3. Improve serialization abstraction - Remove redundant UTF-8 encoding in client - Let transport layer handle all serialization - Pass Python strings directly instead of bytes Signed-off-by: maobaolong <maobaolong@tencent.com> Signed-off-by: baoloongmao <baoloongmao@tencent.com> --------- Signed-off-by: baoloongmao <baoloongmao@tencent.com> Signed-off-by: maobaolong <maobaolong@tencent.com>

maobaolong requested a review from sammshen February 17, 2026 05:36

gemini-code-assist Bot reviewed Feb 17, 2026

View reviewed changes

Comment thread lmcache/v1/lookup_client/lmcache_lookup_client.py Outdated

Comment thread lmcache/v1/rpc/zmq_transport.py Outdated

Comment thread lmcache/v1/lookup_client/lmcache_lookup_client.py

maobaolong added the full Run comprehensive tests on this PR label Feb 19, 2026

sammshen approved these changes Feb 22, 2026

View reviewed changes

maobaolong enabled auto-merge (squash) February 22, 2026 03:08

chunxiaozheng approved these changes Feb 24, 2026

View reviewed changes

sammshen mentioned this pull request Feb 24, 2026

[Roadmap] LMCache Roadmap for 2026 Q1 #2350

Open

24 tasks

maobaolong added 2 commits March 9, 2026 17:34

Refactor lookup client/server and abstract rpc layer.

c9ea275

Signed-off-by: baoloongmao <baoloongmao@tencent.com>

maobaolong force-pushed the ipc_refactor branch from dd558cc to 5c1d69b Compare March 9, 2026 09:37

maobaolong merged commit fa5c5ac into LMCache:dev Mar 9, 2026
24 of 28 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Refactor lookup client/server and abstract rpc layer.#2609

Refactor lookup client/server and abstract rpc layer.#2609
maobaolong merged 2 commits intoLMCache:devfrom
maobaolong:ipc_refactor

maobaolong commented Feb 17, 2026 •

edited

Loading

Uh oh!

gemini-code-assist Bot commented Feb 17, 2026

Uh oh!

gemini-code-assist Bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

maobaolong commented Feb 19, 2026

Uh oh!

sammshen left a comment

Uh oh!

chunxiaozheng left a comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

maobaolong commented Feb 17, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

gemini-code-assist Bot commented Feb 17, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

maobaolong commented Feb 19, 2026

Uh oh!

sammshen left a comment

Choose a reason for hiding this comment

Uh oh!

chunxiaozheng left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

maobaolong commented Feb 17, 2026 •

edited

Loading