misc: Refactor TaskManager init and clean up QueryContextManager

[duxiao1212]

Summary: as title

Differential Revision: D87955845

[sourcery-ai]

Reviewer's Guide

Refactors TaskManager to accept an optional externally-provided QueryContextManager and tightens QueryContextCache implementation details, including explicit constructor, safer initialization, and a corrected LRU eviction loop.

Class diagram for TaskManager and QueryContextManager refactor

classDiagram
  class TaskManager {
    +TaskManager(folly::Executor* driverExecutor, folly::Executor* httpSrvExecutor, folly::Executor* spillerExecutor, std::unique_ptr<QueryContextManager> queryContextManager)
    -std::unique_ptr<QueryContextManager> queryContextManager_
    -velox::exec::OutputBufferManager& bufferManager_
    -folly::Executor* httpSrvCpuExecutor_
    -uint64_t lastNotOverloadedTimeInSecs_
    #void shutdown()
  }

  class QueryContextManager {
    +QueryContextManager(folly::Executor* driverExecutor, folly::Executor* spillerExecutor)
    -QueryContextCache queryContextCache_
  }

  class QueryContextCache {
    +explicit QueryContextCache(size_t initial_capacity)
    +size_t capacity() const
    +std::shared_ptr<velox::core::QueryCtx> insert(protocol::QueryId queryId, std::shared_ptr<velox::core::QueryCtx> queryCtx)
    +void setActive(protocol::QueryId queryId)
    +void setTasksStarted(protocol::QueryId queryId)
    -void evict()
    -size_t capacity_
    -std::list<protocol::QueryId> queryIds_
    -QueryCtxMap queryCtxs_
  }

  class QueryCtxCacheValue {
    +std::weak_ptr<velox::core::QueryCtx> queryCtx
    +std::list<protocol::QueryId>::iterator idListIterator
    +bool hasStartedTasks
  }

  TaskManager --> QueryContextManager : owns
  QueryContextManager --> QueryContextCache : owns
  QueryContextCache --> "*" QueryCtxCacheValue : stores

Loading

Flow diagram for TaskManager constructor QueryContextManager injection

flowchart TD
  A["TaskManager constructor called"] --> B{"queryContextManager parameter is null"}
  B -- "Yes" --> C["Create new QueryContextManager(driverExecutor, spillerExecutor)"]
  B -- "No" --> D["Move provided queryContextManager into queryContextManager_"]
  C --> E["Assign new QueryContextManager to queryContextManager_"]
  D --> E
  E --> F["Initialize bufferManager_ from OutputBufferManager::getInstanceRef()"]
  F --> G["Store httpSrvCpuExecutor_ and initialize lastNotOverloadedTimeInSecs_"]
  G --> H["TaskManager construction complete"]

Loading

Flow diagram for updated QueryContextCache LRU eviction

flowchart TD
  A["evict() called"] --> B["Set victim to queryIds_.rbegin()"]
  B --> C{"victim != queryIds_.rend()"}
  C -- "No" --> H["No eviction performed"]
  C -- "Yes" --> D["Find iter = queryCtxs_.find(*victim)"]
  D --> E{"iter != queryCtxs_.end() and iter->second.queryCtx.lock() is empty"}
  E -- "Yes" --> F["Erase iter from queryCtxs_"]
  F --> G["Erase std::next(victim).base() from queryIds_"]
  G --> I["Return from evict()"]
  E -- "No" --> J["Increment victim (move towards rend)"]
  J --> C

Loading

File-Level Changes

Change	Details	Files
Allow TaskManager to be constructed with an externally-owned QueryContextManager instance.	Extend TaskManager constructor to take a std::unique_ptr with a default of nullptr. Initialize queryContextManager_ with either the provided instance or a new QueryContextManager constructed from driverExecutor and spillerExecutor. Expose shutdown() as a protected method instead of public to better encapsulate lifecycle handling.	presto-native-execution/presto_cpp/main/TaskManager.cpp presto-native-execution/presto_cpp/main/TaskManager.h
Make QueryContextCache construction and insertion more explicit and modernize the eviction logic for better safety and clarity.	Mark QueryContextCache(size_t) as explicit to avoid unintended implicit conversions. Use designated-member-style initialization when inserting entries into queryCtxs_ to bind queryId, weak_ptr, list iterator, and hasStartedTasks flag in a single insert call. Rewrite the eviction loop to iterate with a reverse iterator, look up entries with find() instead of operator[], and erase from both the map and list using valid iterators to avoid undefined behavior. Reorder includes so folly/executors/IOThreadPoolExecutor.h is included after project headers for consistency.	presto-native-execution/presto_cpp/main/QueryContextManager.cpp presto-native-execution/presto_cpp/main/QueryContextManager.h

---

Tips and commands

Interacting with Sourcery

• Trigger a new review: Comment @sourcery-ai review on the pull request.
• Continue discussions: Reply directly to Sourcery's review comments.
• Generate a GitHub issue from a review comment: Ask Sourcery to create an
  issue from a review comment by replying to it. You can also reply to a
  review comment with @sourcery-ai issue to create an issue from it.
• Generate a pull request title: Write @sourcery-ai anywhere in the pull
  request title to generate a title at any time. You can also comment
  @sourcery-ai title on the pull request to (re-)generate the title at any time.
• Generate a pull request summary: Write @sourcery-ai summary anywhere in
  the pull request body to generate a PR summary at any time exactly where you
  want it. You can also comment @sourcery-ai summary on the pull request to
  (re-)generate the summary at any time.
• Generate reviewer's guide: Comment @sourcery-ai guide on the pull
  request to (re-)generate the reviewer's guide at any time.
• Resolve all Sourcery comments: Comment @sourcery-ai resolve on the
  pull request to resolve all Sourcery comments. Useful if you've already
  addressed all the comments and don't want to see them anymore.
• Dismiss all Sourcery reviews: Comment @sourcery-ai dismiss on the pull
  request to dismiss all existing Sourcery reviews. Especially useful if you
  want to start fresh with a new review - don't forget to comment
  @sourcery-ai review to trigger a new review!

Customizing Your Experience

Access your dashboard to:

• Enable or disable review features such as the Sourcery-generated pull request
  summary, the reviewer's guide, and others.
• Change the review language.
• Add, remove or edit custom review instructions.
• Adjust other review settings.

Getting Help

• Contact our support team for questions or feedback.
• Visit our documentation for detailed guides and information.
• Keep in touch with the Sourcery team by following us on X/Twitter, LinkedIn or GitHub.

[sourcery-ai]

Hey there - I've reviewed your changes and found some issues that need to be addressed.

Prompt for AI Agents

Please address the comments from this code review:

## Individual Comments

### Comment 1
<location> `presto-native-execution/presto_cpp/main/TaskManager.h:40-39` </location>
<code_context>
   virtual ~TaskManager() = default;

-  /// Invoked by Presto server shutdown to wait for all the tasks to complete
+ protected:
+
+ /// Invoked by Presto server shutdown to wait for all the tasks to complete
   /// and cleanup the completed tasks.
   void shutdown();
</code_context>

<issue_to_address>
**issue (bug_risk):** Changing the access specifier here may unintentionally make all following TaskManager methods protected.

Placing `protected:` here changes the visibility of `shutdown()` and all following members until the next access specifier. If no later `public:` is declared, this will hide the existing TaskManager public API and change the class interface. If you only intend `shutdown()` to be protected, move `protected:` immediately before `shutdown()` and add `public:` again afterward (or adjust only `shutdown()`’s declaration).
</issue_to_address>

---

Sourcery is free for open source - if you like our reviews please consider sharing them ✨

• X
• Mastodon
• LinkedIn
• Facebook

_{Help me be more useful! Please click 👍 or 👎 on each comment and I'll use the feedback to improve your reviews.}

[sourcery-ai]

issue (bug_risk): Changing the access specifier here may unintentionally make all following TaskManager methods protected.

Placing protected: here changes the visibility of shutdown() and all following members until the next access specifier. If no later public: is declared, this will hide the existing TaskManager public API and change the class interface. If you only intend shutdown() to be protected, move protected: immediately before shutdown() and add public: again afterward (or adjust only shutdown()’s declaration).