misc: Expose TotalRunningTaskCount on Coordinator

[shelton408]

Summary: Expose running task count so we can use it to track how many in flight tasks a coordinator is handling. Due to recent coordinator overloads/crashes, we want to use this metric to do load testing on the coordinator.

Differential Revision: D88803334

== NO RELEASE NOTE ==

[sourcery-ai]

Reviewer's guide (collapsed on small PRs)

Reviewer's Guide

Expose the internal total running task count on InternalResourceGroupManager as a JMX/ODS-managed metric so it can be tracked externally for load testing and monitoring.

Sequence diagram for exposing TotalRunningTaskCount via JMX/ODS

sequenceDiagram
    participant ODSCollector
    participant JMXServer
    participant InternalResourceGroupManager
    participant RootInternalResourceGroup

    ODSCollector->>JMXServer: query metric InternalResourceGroupManager.TotalRunningTaskCount
    JMXServer->>InternalResourceGroupManager: invoke getTotalRunningTaskCount()
    activate InternalResourceGroupManager
    loop for each rootGroup
        InternalResourceGroupManager->>RootInternalResourceGroup: getRunningTaskCount()
        RootInternalResourceGroup-->>InternalResourceGroupManager: runningTaskCount
    end
    InternalResourceGroupManager-->>JMXServer: totalRunningTaskCount
    deactivate InternalResourceGroupManager
    JMXServer-->>ODSCollector: totalRunningTaskCount

Loading

Class diagram for updated InternalResourceGroupManager metric exposure

classDiagram
    class InternalResourceGroupManager {
        - AtomicBoolean taskLimitExceeded
        - List rootGroups
        + int getTaskLimitExceeded()
        + int getTotalRunningTaskCount()
    }

    class RootInternalResourceGroup {
        + int getRunningTaskCount()
    }

    InternalResourceGroupManager --> RootInternalResourceGroup : aggregates

    class Managed {
    }

    Managed <.. InternalResourceGroupManager : annotation on getTotalRunningTaskCount

Loading

File-Level Changes

Change	Details	Files
Expose total running task count from InternalResourceGroupManager as a managed metric.	Change getTotalRunningTaskCount from private to public so it can be accessed externally. Annotate getTotalRunningTaskCount with @Managed to export it as a JMX/ODS metric. Keep the existing aggregation logic over rootGroups intact for computing the running task count.	presto-main-base/src/main/java/com/facebook/presto/execution/resourceGroups/InternalResourceGroupManager.java

---

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 - here's some feedback:

• Exposing getTotalRunningTaskCount via @Managed means it will be called periodically by the metrics system; consider whether iterating over all rootGroups is cheap enough or if you should cache or otherwise bound the cost of this computation.
• Now that getTotalRunningTaskCount is public, verify that access to rootGroups is thread-safe for concurrent metric reads; if not, consider synchronization or using a thread-safe collection to avoid races.

Prompt for AI Agents

Please address the comments from this code review:

## Overall Comments
- Exposing `getTotalRunningTaskCount` via `@Managed` means it will be called periodically by the metrics system; consider whether iterating over all `rootGroups` is cheap enough or if you should cache or otherwise bound the cost of this computation.
- Now that `getTotalRunningTaskCount` is public, verify that access to `rootGroups` is thread-safe for concurrent metric reads; if not, consider synchronization or using a thread-safe collection to avoid races.

---

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.}

[spershin]

@shelton408
TotalRunningTaskCount is now possible to get from the workers.
Do we need that from the coord?

[tdcmeehan]

Can you use SqlQueryManager.getRunningTaskCount()?