Skip to content

Initialize OOM protection in GrpcQueryServer#15271

Open
jitendrakr88 wants to merge 1 commit intoapache:masterfrom
jitendrakr88:oom-grpc-regression2
Open

Initialize OOM protection in GrpcQueryServer#15271
jitendrakr88 wants to merge 1 commit intoapache:masterfrom
jitendrakr88:oom-grpc-regression2

Conversation

@jitendrakr88
Copy link
Contributor

@jitendrakr88 jitendrakr88 commented Mar 14, 2025

Description

When OOM protection is enabled, the grpc streming queries always return 0 rows. This issue has been observered with spark connector as well as trino connector (reported by another user).

from pyspark.sql.functions import col

cutoff_seconds_since_epoch = 1741169418

df = spark \
    .read \
    .format("pinot") \
    .option("controller", "<host>:<port>") \
    .option("table", "<table>") \
    .option("useGrpcServer", "true") \
    .option("tableType", "REALTIME") \
    .option("segmentsPerSplit", 1) \
    .load() \
    .filter(col("secondsSinceEpoch") >= cutoff_seconds_since_epoch)

print("Number of rows: " + str(df.count())) 
image

Observation:

  • Number of rows returned is always zero if OOM is enabled on server tenant
  • After disabling OOM, the grpc streaming query works fine.
  • After disabling useGrpcServer option, the query works fine with http transport.

Solution

This PR fixes the issue by initialising the missing tracing context in GrpcQueryServer for OOM protection.

_grpcQueryServer.submit(serverRequest, _responseObserver);

// Assert: ThreadAccountantOps.setupRunner() was called once.
mockedStatic.verify(() -> Tracing.ThreadAccountantOps.setupRunner(anyString()), times(1));
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This test is fine, but can we write a test case which actually validates the gRPC query server can execute queries?

The reason I'm insistent on the test is that the Grpc transport is secondary to the the HttpServer in terms of importance and when folks do sweeping changes (like the tracing setup), regressions to the gRPC server slip unnoticed.

import static org.mockito.Mockito.*;

public class GrpcQueryServerTest {
private GrpcQueryServer _grpcQueryServer;
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

(minor) Please fix the format

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

3 participants