Eliminate per-segment DistinctTable allocation and queue contention in DistinctCombineOperator - #19092
Open
Akanksha-kedia wants to merge 11 commits into
Open
Conversation
…validation Implements a Trino/Ambari-style session layer for the Pinot controller UI that works on top of any configured AccessControlFactory (BasicAuth, ZkBasicAuth, LDAP, or custom). **Security problems fixed:** - Credentials were visible in every browser network tab request (Authorization header) - Logout only cleared localStorage — server session survived indefinitely - No inactivity timeout; sessions never expired - Credentials stored in JS memory (XSS-readable) - Broker auth forwarded directly from the browser - Race condition in session lookup (non-atomic get→check→remove) - No CSRF protection on session cookie **How it works:** 1. `POST /auth/login` validates credentials via the configured AccessControl, creates a server-side session, and sets an HttpOnly + SameSite=Strict cookie 2. `GET /auth/logout` immediately removes the session from the server store (proper invalidation) 3. `GET /auth/session` checks session validity for page-refresh restoration 4. `SessionAuthenticationFilter` (AUTHENTICATION-10 priority) validates the cookie before `AuthenticationFilter` runs; passes through requests with Authorization header for API compat 5. `PinotQueryResource` injects stored Basic-auth token server-side when forwarding to broker **New files:** - `SessionManager.java`: ConcurrentHashMap session store with sliding TTL, atomic compute() - `SessionAuthenticationFilter.java`: JAX-RS filter for cookie-based auth - `SessionBasicAuthAccessControlFactory.java`: convenience wrapper factory - `PinotSessionLoginResource.java`: /auth/login, /auth/logout, /auth/session endpoints - `SessionManagerTest.java`: unit + 50-thread concurrency test **Configuration (all default-off for backward compatibility):** ```properties controller.ui.session.authentication.enabled=true controller.ui.session.cookie.secure=true # set false for HTTP-only dev controller.ui.session.inactivity.timeout.seconds=300 # default 5 min ``` **Frontend changes:** - `Models.ts`: SESSION enum added to AuthWorkflow - `axios-config.ts`: SESSION workflow deletes Authorization header (cookie sent automatically) - `AuthProvider.tsx`: reads timeout from /auth/info, checks /auth/session on page refresh - `App.tsx`: inactivity timer with 60-second warning dialog before auto-logout - `Header.tsx`: real logout button that calls /auth/logout server-side - `LoginPage.tsx`: SESSION path POSTs form data to /auth/login (no Basic header in JS)
- Remove trailing whitespace on blank lines in SessionAuthenticationFilter - Remove blank lines before end-of-file in 4 new/modified files - Replace Collections.emptyList/emptyMap with List.of/Map.of in PinotSessionLoginResource - Fix LeftCurly violation in SessionManagerTest (synchronized block brace) Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
- SessionBasicAuthAccessControlFactory: replace non-existent BasicAuthUtils/ AbstractBasicAuthAccessControl with correct BasicAuthPrincipalUtils and inline AccessControl implementation (mirrors BasicAuthAccessControlFactory) - PinotQueryResource: fix @optional annotation package from org.glassfish.hk2.api to org.jvnet.hk2.annotations (correct HK2 location) Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
…return AuthenticationFilter: replace early `return` with Basic-auth token injection from the session store into the request headers, so validatePermission() and FineGrainedAuthUtils still run with the session user's identity. Extract isSessionEnabled() helper to consolidate the repeated config check. PinotQueryResource: remove `if (!sessionValid)` guards around hasAccess() calls in getQueryResponse() and getMultiStageQueryResponse(). Now that AuthenticationFilter injects the Authorization header for session requests, the downstream hasAccess() calls receive the session user's credentials and work correctly. Remove hasValidSessionCookie() and sessionValid parameter. ControllerConf: shorten multi-line session constant Javadocs to one-liners. Addresses CRITICAL review feedback from @xiangfu0 and @shounakmk219 on PR apache#18933. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
…roker auth ControllerAdminApiApplication: revert CORS to wildcard-only (no credentials). The Pinot UI is same-origin so CORS is not needed; the previous origin-reflect + Allow-Credentials approach was too permissive for same-site sibling origins. PinotSessionLoginResource: use hasAccess(AccessType, headers, url) instead of hasAccess(null, AccessType, headers, url). The table-scoped overload calls hasTable(null) which rejects principals with a table allowlist, locking them out of session login entirely. The identity-only overload validates credentials without requiring any table permission. PinotQueryResource: extract injectSessionCredentialIfNeeded() helper and call it in both sendRequestToBroker() and sendTimeSeriesRequestToBroker(). Previously only the SQL forward path re-injected the session credential after extractHeaders() stripped it; time-series requests were forwarded without broker auth in session mode. Addresses review feedback from @xiangfu0 on PR apache#18933. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
…fault impl Extract SessionManager as an interface (keeping SESSION_COOKIE_NAME and DEFAULT_SESSION_TTL_SECONDS constants + the public API) and rename the in-memory implementation to InMemorySessionManager. All injection points (AuthenticationFilter, SessionAuthenticationFilter, PinotQueryResource, PinotSessionLoginResource) already reference the SessionManager type so they are unaffected. BaseControllerStarter now instantiates InMemorySessionManager. SessionManagerTest updated to construct InMemorySessionManager directly while keeping variable types as SessionManager. Addresses @shounakmk219 review feedback on PR apache#18933. Co-Authored-By: Claude Sonnet 4.6 (1M context) <noreply@anthropic.com>
… API clients extractHeaders was stripping the Authorization header from every request when session mode was globally enabled, including API clients that supply their own Basic/Bearer token with no session cookie. injectSessionCredentialIfNeeded only re-injected if a session cookie was present, so API clients had their auth silently removed before the broker call. Merge the two-step strip+inject into a single buildBrokerHeaders call that resolves the session credential exactly once via resolveSessionCredential. If a valid session cookie is found, the caller Authorization is replaced with the stored token; otherwise all headers are forwarded unchanged.
…ticky-session requirement Three issues raised in review: 1. TTL mismatch: InMemorySessionManager used a sliding-window TTL (getUsername extended expiry on each call) while the browser cookie Max-Age was fixed at login. An active session could outlive its browser cookie, causing unexpected logouts. Fixed by removing the sliding extension in getUsername; both server session and cookie now expire at loginTime + TTL. 2. Cross-controller logout: InMemorySessionManager is per-process, so logout routed to a different controller in an LB setup silently fails. Added a comment in BaseControllerStarter documenting the sticky-session requirement and the degraded-but-functional behavior without it. 3. CORS allowed headers: Access-Control-Allow-Headers was missing the custom Pinot 'Database' header used by DatabaseUtils for multi-database table resolution. Also adds If-Match for segment-upload endpoints. Browser clients that set these headers were failing preflight.
…actory, escapeJson, timer leak - AuthenticationFilter: also detect SESSION workflow from factory.getAuthWorkflowInfo() so isSessionEnabled() works even when CONTROLLER_UI_SESSION_ENABLED flag is absent. Cache result in volatile Boolean to avoid calling factory.create() on every request. Add auth/session/renew to UNPROTECTED_PATHS so the renew endpoint is reachable without an Authorization header. - PinotSessionLoginResource: replace non-atomic get/invalidate/create with rotateSession() to eliminate TOCTOU in POST /auth/session/renew. Add escapeJson() and apply to all three username-in-JSON response sites. Fix buildDeleteCookieHeader to propagate the Secure flag. - SessionManager: add rotateSession() default method (non-atomic fallback for custom impls). - InMemorySessionManager: override rotateSession() with CAS-based atomic rotation using ConcurrentHashMap.remove(key, value) to prevent double-rotation under concurrent requests. - PinotControllerAuthResource: enrich AuthWorkflowInfo with inactivity timeout when factory advertises SESSION workflow, not just when the explicit config flag is set. - App.tsx: change "Stay Logged In" to POST /auth/session/renew (was incorrectly hitting GET /auth/session which doesn't renew server-side TTL). Fix performSessionLogout to also clear the countdown interval, preventing warningCountdown going negative after logout.
…n DistinctCombineOperator Replace the one-result-block-per-segment / BlockingQueue pattern with per-task thread-local DistinctTable accumulation, following the same design as GroupByCombineOperator: - Each worker task is assigned a unique slot (0..numTasks-1) and merges every segment it processes directly into that slot's DistinctTable, eliminating both the per-segment DistinctTable allocation and the LinkedBlockingQueue offer/poll overhead. - After all tasks finish (via CountDownLatch), the main thread merges at most numTasks tables instead of numSegments tables. - Early termination sets _nextOperatorId to _numOperators to stop all tasks as soon as the LIMIT is satisfied, consistent with the existing pattern. - Adds DistinctCombineOperatorTest covering full-scan, early termination, ORDER BY ASC/DESC, and cross-segment deduplication.
Contributor
Author
|
cc @Jackie-Jiang @walterddr @richardstartin — would appreciate a review when you get a chance! |
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #19092 +/- ##
=============================================
- Coverage 65.05% 37.84% -27.22%
- Complexity 1403 1420 +17
=============================================
Files 3399 3435 +36
Lines 212808 218354 +5546
Branches 33568 34706 +1138
=============================================
- Hits 138443 82626 -55817
- Misses 63223 128128 +64905
+ Partials 11142 7600 -3542
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Harness. 🚀 New features to boost your workflow:
|
Contributor
Author
|
Could a maintainer please add the |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
DistinctCombineOperatorpreviously followed the genericBaseSingleBlockCombineOperatorqueue-based pattern: each segment operator posted aDistinctResultsBlockto aLinkedBlockingQueue, and the main thread serially dequeued and merged them — oneDistinctTableallocation and one queue round-trip per segment.This PR replaces that with per-task thread-local accumulation, following the same design already used by
GroupByCombineOperator:Changes
DistinctCombineOperator[0, numTasks)viaAtomicIntegerand merges every segment it processes directly into that slot'sDistinctTable— no queue, no per-segment allocation.CountDownLatch), the main thread merges at mostnumTasksper-task tables rather thannumSegmentstables._nextOperatorIdto_numOperatorsas soon as a task's accumulated table satisfies the LIMIT, stopping all tasks.CountDownLatch.countDown()→await()establishes the happens-before relationship that makes per-task table writes visible to the main thread (same guarantee used byGroupByCombineOperator).DistinctCombineOperatorTest(new)numSegments × uniquePerSegmentdistinct values collectedPerformance impact
For a
SELECT DISTINCT col LIMIT Lquery over S segments with T worker threads (T << S):DistinctTableallocations + S queue operations + S serial mergesDistinctTableallocations + T final mergesThe benefit is most pronounced for queries with a small
LIMITover many segments, where the bounded table size makes allocation cost non-trivial.Test plan
DistinctCombineOperatorTest— 5 unit tests, all passingDistinctQueriesTestintegration tests pass