Skip to content

Avoid data race in grpc alarm (#10124) #11030

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
ti-chi-bot wants to merge 1 commit into
base: release-8.5
Choose a base branch
from
+21 −4
Open

Avoid data race in grpc alarm (#10124) #11030

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions dbms/src/Flash/EstablishCall.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -180,6 +180,11 @@ void EstablishCallData::initRpc()
}
}

grpc::Alarm & EstablishCallData::getAlarm()
{
return alarm;
}

void EstablishCallData::tryConnectTunnel()
{
auto * task_manager = service->getContext()->getTMTContext().getMPPTaskManager().get();
Expand Down
4 changes: 4 additions & 0 deletions dbms/src/Flash/EstablishCall.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,6 +19,7 @@
#include <Common/Stopwatch.h>
#include <Flash/FlashService.h>
#include <Flash/Mpp/MPPTaskId.h>
#include <grpcpp/alarm.h>
#include <kvproto/tikvpb.grpc.pb.h>

namespace DB
Expand Down Expand Up @@ -81,6 +82,8 @@ class EstablishCallData final
grpc::ServerContext * getGrpcContext() { return &ctx; }

String getResourceGroupName() const { return resource_group_name; }
grpc::Alarm & getAlarm();


private:
/// WARNING: Since a event from one grpc completion queue may be handled by different
Expand Down Expand Up @@ -138,5 +141,6 @@ class EstablishCallData final
String resource_group_name;
String connection_id;
double waiting_task_time_ms = 0;
grpc::Alarm alarm{};
};
} // namespace DB
13 changes: 10 additions & 3 deletions dbms/src/Flash/Mpp/MPPTaskManager.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -56,7 +56,7 @@ void MPPGatherTaskSet::cancelAlarmsBySenderTaskId(const MPPTaskId & task_id)
if (alarm_it != alarms.end())
{
for (auto & alarm : alarm_it->second)
alarm.second.Cancel();
alarm.second.get().Cancel();
alarms.erase(alarm_it);
}
}
Expand Down Expand Up @@ -247,8 +247,16 @@ std::pair<MPPTunnelPtr, String> MPPTaskManager::findAsyncTunnel(
context.getSettingsRef().auto_spill_check_min_interval_ms.get());
if (gather_task_set == nullptr)
gather_task_set = query->addMPPGatherTaskSet(id.gather_id);
<<<<<<< HEAD
auto & alarm = gather_task_set->alarms[sender_task_id][receiver_task_id];
call_data->setToWaitingTunnelState();
=======
auto & alarm = call_data->getAlarm();
call_data->setCallStateAndUpdateMetrics(
EstablishCallData::WAIT_TUNNEL,
GET_METRIC(tiflash_establish_calldata_count, type_wait_tunnel_calldata));
gather_task_set->alarms[sender_task_id].emplace(receiver_task_id, std::ref(alarm));

@coderabbitai coderabbitai Bot Aug 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

🩺 Stability & Availability | 🔴 Critical | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

rg -n -C 8 \
  'EstablishMPPConnection|findAsyncTunnel|sender_meta\(\).*task_id|receiver_meta\(\).*task_id' .

Repository: pingcap/tiflash

Length of output: 50026

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "== MPPTaskManager relevant sections =="
sed -n '170,255p;360,395p' dbms/src/Flash/Mpp/MPPTaskManager.cpp

echo
echo "== MPPTaskManager.h gather_task_set / alarm declarations =="
sed -n '180,230p' dbms/src/Flash/Mpp/MPPTaskManager.h

echo
echo "== Find MPPGatherTaskSet definition and alarm methods =="
rg -n -C 4 'struct MPPGatherTaskSet|class MPPGatherTaskSet|alarms|hasAlarm|notifyAlarm|Cancel\\(' dbms/src/Flash/Mpp

Repository: pingcap/tiflash

Length of output: 6992

🏁 Script executed:

#!/bin/bash
set -euo pipefail

echo "== MPPTaskManager.h declarations =="
sed -n '1,140p' dbms/src/Flash/Mpp/MPPTaskManager.h
sed -n '140,190p' dbms/src/Flash/Mpp/MPPTaskManager.h
sed -n '190,270p' dbms/src/Flash/Mpp/MPPTaskManager.h

echo
echo "== findAsyncTunnel full and concurrent merge markers =="
sed -n '204,300p' dbms/src/Flash/Mpp/MPPTaskManager.cpp

echo
echo "== alarm-related symbols and Cancel references =="
rg -n -C 3 'alarms|hasAlarm|notifyAlarm|setAlarm|Cancel\(' dbms/src/Flash dbms/src/Core 2>/dev/null || true

echo
echo "== compile marker check =="
python3 - <<'PY'
from pathlib import Path
p=Path('dbms/src/Flash/Mpp/MPPTaskManager.cpp')
for i,line in enumerate(p.read_text().splitlines(),1):
    if line.startswith('<'):
        print(i, line)
PY

Repository: pingcap/tiflash

Length of output: 25014

Resolve the merge conflict in findAsyncTunnel().

dbms/src/Flash/Mpp/MPPTaskManager.cpp:250-259 still contains unresolved conflict markers, so this translation unit will not compile.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@dbms/src/Flash/Mpp/MPPTaskManager.cpp` at line 258, Resolve the remaining
merge-conflict markers in findAsyncTunnel(), preserving the intended
gather_task_set->alarms insertion involving sender_task_id, receiver_task_id,
and alarm. Remove all conflict markers and ensure the resulting C++ block is
syntactically valid and compilable.

>>>>>>> 262b942077 (Avoid data race in grpc alarm (#10124))
Comment on lines +250 to +259

@coderabbitai coderabbitai Bot Aug 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

🎯 Functional Correctness | 🔴 Critical | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

file="dbms/src/Flash/Mpp/MPPTaskManager.cpp"

if rg -n '^(<<<<<<<|=======|>>>>>>>)' "$file"; then
  echo "Unresolved merge conflict markers found" >&2
  exit 1
fi

rg -n -C 3 \
  'setCallStateAndUpdateMetrics|WAIT_TUNNEL|setToWaitingTunnelState|getAlarm|alarms.*emplace' \
  dbms/src/Flash/EstablishCall.h \
  "$file"

Repository: pingcap/tiflash

Length of output: 286

🏁 Script executed:

#!/bin/bash
set -u

echo "== file section =="
sed -n '230,280p' dbms/src/Flash/Mpp/MPPTaskManager.cpp 2>/dev/null || true

echo "== EstablishCall.h related section =="
sed -n '1,220p' dbms/src/Flash/EstablishCall.h 2>/dev/null || true

echo "== other relevant files =="
git ls-files | rg 'dbms/src/Flash/.*\.(h|cpp|hpp)$' | sed -n '1,120p'

echo "== search for API names =="
rg -n 'setCallStateAndUpdateMetrics|WAIT_TUNNEL|setToWaitingTunnelState|getAlarm|type_wait_tunnel_calldata|tiflash_establish_calldata_count' dbms/src/Flash 2>/dev/null || true

Repository: pingcap/tiflash

Length of output: 14831

🏁 Script executed:

#!/bin/bash
set -u

echo "== MPPTaskManager outline == "
ast-grep outline dbms/src/Flash/Mpp/MPPTaskManager.cpp --view compact 2>/dev/null || true

echo "== alarms declarations and usages in MPPTaskManager.cpp =="
rg -n -C 2 'gather_task_set->alarms|sender_task_id|receiver_task_id|\.Set\(|std::ref\(alarm\)|type_wait_tunnel_calldata' dbms/src/Flash/Mpp/MPPTaskManager.cpp

echo "== EstablishCall related symbols =="
rg -n -C 2 'WAIT_TUNNEL|getAlarm|setToWaitingTunnelState|setCallStateAndUpdateMetrics|type_wait_tunnel_calldata|tiflash_establish_calldata_count|GET_METRIC' dbms/src/Flash/EstablishCall.h dbms/src/Flash/EstablishCall.cpp dbms/src/Metrics 2>/dev/null || true

echo "== changed conflict metadata =="
git diff --stat || true
git diff --name-status || true

Repository: pingcap/tiflash

Length of output: 6465

Resolve the unresolved merge conflict at MPPTaskManager.cpp:250-259.

The conflict markers remain in the build file, so MPPTaskManager.cpp does not compile. The incoming branch also calls setCallStateAndUpdateMetrics() and EstablishCallData::WAIT_TUNNEL, but dbms/src/Flash/EstablishCall.h exposes setToWaitingTunnelState() and keeps WAIT_TUNNEL private; getAlarm() is available, so keep std::ref(alarm) as intended and update the public API instead.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@dbms/src/Flash/Mpp/MPPTaskManager.cpp` around lines 250 - 259, Resolve the
conflict in the MPP task setup by removing all merge markers and retaining the
race-safe alarm retrieval via EstablishCallData::getAlarm() and std::ref(alarm).
Update EstablishCallData’s public API in EstablishCall.h so the call uses the
existing public setToWaitingTunnelState() behavior, or otherwise exposes an
equivalent valid transition without referencing private WAIT_TUNNEL or
unavailable setCallStateAndUpdateMetrics().

if likely (cq != nullptr)
{
alarm.Set(cq, Clock::now() + std::chrono::seconds(10), call_data->asGRPCKickTag());
Expand Down Expand Up @@ -280,7 +288,6 @@ std::pair<MPPTunnelPtr, String> MPPTaskManager::findAsyncTunnel(
}
}
/// don't need to delete the alarm here because registerMPPTask will delete all the related alarm

return task->getTunnel(request);
}

Expand Down Expand Up @@ -375,7 +382,7 @@ void MPPTaskManager::abortMPPGather(const MPPGatherId & gather_id, const String
for (auto & alarms_per_task : gather_task_set->alarms)
{
for (auto & alarm : alarms_per_task.second)
alarm.second.Cancel();
alarm.second.get().Cancel();
}
gather_task_set->alarms.clear();
if (!gather_task_set->hasMPPTask())
Expand Down
3 changes: 2 additions & 1 deletion dbms/src/Flash/Mpp/MPPTaskManager.h
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,7 @@

#include <chrono>
#include <condition_variable>
#include <functional>
#include <memory>
#include <mutex>
#include <unordered_map>
Expand All @@ -42,7 +43,7 @@ struct MPPGatherTaskSet
State state = Normal;
String error_message;
/// <sender_task_id, <receiver_task_id, alarm>>
std::unordered_map<Int64, std::unordered_map<Int64, grpc::Alarm>> alarms;
std::unordered_map<Int64, std::unordered_map<Int64, std::reference_wrapper<grpc::Alarm>>> alarms;

@coderabbitai coderabbitai Bot Aug 10, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

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

🩺 Stability & Availability | 🟠 Major | ⚡ Quick win

🧩 Analysis chain

🏁 Script executed:

#!/bin/bash
set -euo pipefail

rg -n -C 8 \
  'delete this|EstablishCallData::~|getAlarm\(\)|alarms|cancelAlarmsBySenderTaskId|abortMPPGather|Cancel\(\)' \
  dbms/src/Flash/EstablishCall.h \
  dbms/src/Flash/EstablishCall.cpp \
  dbms/src/Flash/Mpp/MPPTaskManager.h \
  dbms/src/Flash/Mpp/MPPTaskManager.cpp

rg -n -C 4 \
  'ComputeServerRunner_testErrorMessage|testErrorMessage|ThreadSanitizer|TSAN' .

Repository: pingcap/tiflash

Length of output: 31996

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Inspect the relevant class definitions, methods, and exact call sites using small targeted reads.
printf '--- MPPTaskManager.h outline ---\n'
ast-grep outline dbms/src/Flash/Mpp/MPPTaskManager.h --view compact || true

printf '--- MPPTaskManager.h relevant lines ---\n'
sed -n '1,110p' dbms/src/Flash/Mpp/MPPTaskManager.h
sed -n '240,275p' dbms/src/Flash/Mpp/MPPTaskManager.h

printf '--- MPPTaskManager.cpp relevant lines ---\n'
sed -n '35,70p' dbms/src/Flash/Mpp/MPPTaskManager.cpp
sed -n '230,310p' dbms/src/Flash/Mpp/MPPTaskManager.cpp
sed -n '344,405p' dbms/src/Flash/Mpp/MPPTaskManager.cpp
sed -n '488,525p' dbms/src/Flash/Mpp/MPPTaskManager.cpp

printf '--- EstablishCall.h relevant lines ---\n'
sed -n '1,130p' dbms/src/Flash/EstablishCall.h

printf '--- EstablishCall.cpp relevant lines ---\n'
sed -n '62,145p' dbms/src/Flash/EstablishCall.cpp
sed -n '140,190p' dbms/src/Flash/EstablishCall.cpp
sed -n '270,340p' dbms/src/Flash/EstablishCall.cpp

printf '--- Tests relevant lines ---\n'
sed -n '1354,1430p' dbms/src/Flash/tests/gtest_compute_server.cpp

printf '--- Deterministic control-flow extraction ---\n'
python3 - <<'PY'
from pathlib import Path
for path in [
    Path('dbms/src/Flash/Mpp/MPPTaskManager.cpp'),
    Path('dbms/src/Flash/Mpp/MPPTaskManager.h'),
    Path('dbms/src/Flash/EstablishCall.cpp'),
    Path('dbms/src/Flash/EstablishCall.h'),
    Path('dbms/src/Flash/tests/gtest_compute_server.cpp'),
]:
    text = path.read_text()
    print(f'## {path}')
    for needle in [
        'std::reference_wrapper<grpc::Alarm> alarms',
        'cancelAlarmsBySenderTaskId',
        'markTaskAsFinishedOrFailed',
        'abortMPPGather',
        'alarm.second.get().Cancel()',
        'gather_task_set->alarms.clear()',
        'gather_task_set->alarms[sender_task_id].emplace(receiver_task_id, std::ref(alarm))',
        'call_data->getAlarm()',
        'delete this;',
        'tryConnectTunnel()',
        'async_tunnel_sender->setResult',
        'ComputeServerRunner',
        'testErrorMessage',
    ]:
        idx = text.find(needle)
        if idx == -1:
            continue
        line = text[:idx].count('\n') + 1
        print(f'line {line}: {needle}')
PY

Repository: pingcap/tiflash

Length of output: 244

🏁 Script executed:

#!/bin/bash
set -euo pipefail

# Inspect the relevant class definitions, methods, and exact call sites using small targeted reads.
printf '%s\n' '--- MPPTaskManager.h outline ---'
ast-grep outline dbms/src/Flash/Mpp/MPPTaskManager.h --view compact || true

printf '%s\n' '--- MPPTaskManager.h relevant lines ---'
sed -n '1,110p' dbms/src/Flash/Mpp/MPPTaskManager.h
sed -n '240,275p' dbms/src/Flash/Mpp/MPPTaskManager.h

printf '%s\n' '--- MPPTaskManager.cpp relevant lines ---'
sed -n '35,70p' dbms/src/Flash/Mpp/MPPTaskManager.cpp
sed -n '230,310p' dbms/src/Flash/Mpp/MPPTaskManager.cpp
sed -n '344,405p' dbms/src/Flash/Mpp/MPPTaskManager.cpp
sed -n '488,525p' dbms/src/Flash/Mpp/MPPTaskManager.cpp

printf '% s\n' '--- EstablishCall.h relevant lines ---'
sed -n '1,130p' dbms/src/Flash/EstablishCall.h

printf '%s\n' '--- EstablishCall.cpp relevant lines ---'
sed -n '62,145p' dbms/src/Flash/EstablishCall.cpp
sed -n '140,190p' dbms/src/Flash/EstablishCall.cpp
sed -n '270,340p' dbms/src/Flash/EstablishCall.cpp

printf '% s\n' '--- Tests relevant lines ---'
sed -n '1354,1430p' dbms/src/Flash/tests/gtest_compute_server.cpp

printf '%s\n' '--- Deterministic control-flow extraction ---'
python3 - <<'PY'
from pathlib import Path
for path in [
    Path('dbms/src/Flash/Mpp/MPPTaskManager.cpp'),
    Path('dbms/src/Flash/Mpp/MPPTaskManager.h'),
    Path('dbms/src/Flash/EstablishCall.cpp'),
    Path('dbms/src/Flash/EstablishCall.h'),
    Path('dbms/src/Flash/tests/gtest_compute_server.cpp'),
]:
    text = path.read_text()
    print(f'## {path}')
    for needle in [
        'std::reference_wrapper<grpc::Alarm> alarms',
        'cancelAlarmsBySenderTaskId',
        'markTaskAsFinishedOrFailed',
        'abortMPPGather',
        'alarm.second.get().Cancel()',
        'gather_task_set->alarms.clear()',
        'gather_task_set->alarms[sender_task_id].emplace(receiver_task_id, std::ref(alarm))',
        'call_data->getAlarm()',
        'delete this;',
        'tryConnectTunnel()',
        'async_tunnel_sender->setResult',
        'ComputeServerRunner',
        'testErrorMessage',
    ]:
        idx = text.find(needle)
        if idx == -1:
            continue
        line = text[:idx].count('\n') + 1
        print(f'line {line}: {needle}')
PY

Repository: pingcap/tiflash

Length of output: 30566

🏁 Script executed:

#!/bin/bash
set -euo pipesse

printf '%s\n' '--- EstablishCall.h tail members ---'
sed -n '130,185p' dbms/src/Flash/EstablishCall.h

printf '%s\n' '--- EstablishCall.cpp writeDone ---'
sed -n '310,380p' dbms/src/Flash/EstablishCall.cpp

printf '%s\n' '--- MPPTaskManager.cpp getMPPQueryAndGatherTaskSet ---'
sed -n '310,350p' dbms/src/Flash/Mpp/MPPTaskManager.cpp
sed -n '525,580p' dbms/src/Flash/Mpp/MPPTaskManager.cpp

printf '%s\n' '--- MPPTaskManager.h gather definitions ---'
sed -n '53,180p' dbms/src/Flash/Mpp/MPPTaskManager.h

printf '%s\n' '--- cancelGather calls ---'
rg -n -C 6 'cancelGather' .

Repository: pingcap/tiflash

Length of output: 201

🏁 Script executed:

#!/bin/bash
set -euo pipefail

printf '%s\n' '--- EstablishCall.h tail members ---'
sed -n '130,185p' dbms/src/Flash/EstablishCall.h

printf '%s\n' '--- EstablishCall.cpp writeDone ---'
sed -n '310,380p' dbms/src/Flash/EstablishCall.cpp

printf '%s\n' '--- MPPTaskManager.cpp cleanup helpers ---'
sed -n '310,350p' dbms/src/Flash/Mpp/MPPTaskManager.cpp
sed -n '525,580p' dbms/src/Flash/Mpp/MPPTaskManager.cpp

printf '%s\n' '--- MPPTaskManager.h gather definitions ---'
sed -n '53,180p' dbms/src/Flash/Mpp/MPPTaskManager.h

printf '%s\n' '--- cancelGather calls ---'
rg -n -C 6 'cancelGather' .

Repository: pingcap/tiflash

Length of output: 32258

Keep the stored grpc::Alarm references owned while cancelling.

std::reference_wrapper does not extend the EstablishCallData lifetime, and MpPTaskManager::abortMPPGather() cancels and erases the wrappers even while the owned EstablishCallData::alarm can still be used by the completion queue. Store the alarms/owner until the alarm is removed, or add an object lifetime handle such that a wrapper can reach the alarm after abortMPPGather() erases it. Run the ComputeServerRunner_testErrorMessage ThreadSanitizer reproduction after resolving the merge conflict.

🤖 Prompt for AI Agents 
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.

In `@dbms/src/Flash/Mpp/MPPTaskManager.h` at line 46, Update the alarms storage in
MPPTaskManager and the abortMPPGather cancellation flow so each grpc::Alarm
remains owned, along with its EstablishCallData owner, until the completion
queue removes the alarm; do not erase a non-owning reference while the alarm may
still be used. Preserve cancellation behavior and run the
ComputeServerRunner_testErrorMessage ThreadSanitizer reproduction after
resolving the merge conflict.

/// only used in scheduler
std::queue<MPPTaskId> waiting_tasks;
bool isInNormalState() const { return state == Normal; }
Expand Down