Skip to content

Commit b294d8b

Browse files
shrinidhijoshishrinidhijoshi
authored
fix(native-pos): Fix MaterializedOutput correctness and crash bugs in close/finish lifecycle (#27833)
1 parent 0da3bb8 commit b294d8b

4 files changed

Lines changed: 127 additions & 29 deletions

File tree

‎presto-native-execution/presto_cpp/main/operators/MaterializedOutput.cpp‎

Lines changed: 25 additions & 26 deletions
Original file line numberDiff line numberDiff line change
@@ -488,44 +488,43 @@ bool MaterializedOutput::isFinished() {
488488
return finished_;
489489
}
490490

491+
void MaterializedOutput::recordBufferStats() {
492+
for (const auto& [key, value] : buffer_->stats()) {
493+
addRuntimeStat(key, velox::RuntimeCounter(value));
494+
}
495+
}
496+
491497
void MaterializedOutput::close() {
492-
finish();
498+
if (!finished_) {
499+
// If finish() was never called via noMoreInput(), we are on the error
500+
// path. Abort the buffer instead of attempting to flush.
501+
buffer_->abort();
502+
}
503+
recordBufferStats();
493504
Operator::close();
494505
}
495506

496507
void MaterializedOutput::finish() {
497508
if (finished_) {
498509
return;
499510
}
500-
finished_ = true;
511+
501512
flushBatch();
502513

503-
// Use Velox's allPeersFinished barrier — returns true only for the last
504-
// driver to reach this point. The last driver drains and closes the writer.
505-
// Wrap in try/catch because allPeersFinished throws if the task is
506-
// already terminating (e.g., due to an error on another pipeline).
507-
// In that case, the buffer's destructor will handle cleanup via abort().
508-
try {
509-
std::vector<velox::ContinuePromise> promises;
510-
std::vector<std::shared_ptr<velox::exec::Driver>> peers;
511-
velox::ContinueFuture peerFuture;
512-
auto* driverCtx = operatorCtx()->driverCtx();
513-
bool isLast = driverCtx->task->allPeersFinished(
514-
planNodeId(), driverCtx->driver, &peerFuture, promises, peers);
515-
516-
if (isLast) {
517-
buffer_->noMoreData();
518-
for (const auto& [key, value] : buffer_->stats()) {
519-
addRuntimeStat(key, velox::RuntimeCounter(value));
520-
}
521-
for (auto& promise : promises) {
522-
promise.setValue();
523-
}
514+
std::vector<velox::ContinuePromise> peerPromises;
515+
std::vector<std::shared_ptr<velox::exec::Driver>> peers;
516+
velox::ContinueFuture peerFuture;
517+
auto* driverCtx = operatorCtx()->driverCtx();
518+
auto isLast = driverCtx->task->allPeersFinished(
519+
planNodeId(), driverCtx->driver, &peerFuture, peerPromises, peers);
520+
521+
if (isLast) {
522+
buffer_->noMoreData();
523+
for (auto& promise : peerPromises) {
524+
promise.setValue();
524525
}
525-
} catch (const velox::VeloxRuntimeError&) {
526-
// Task is terminating — abort the buffer if not already closed.
527-
buffer_->abort();
528526
}
527+
529528
finished_ = true;
530529
}
531530

‎presto-native-execution/presto_cpp/main/operators/MaterializedOutput.h‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -125,10 +125,12 @@ class MaterializedOutput : public velox::exec::Operator {
125125
// Project input columns and lazy-load all children.
126126
void initializeInput(velox::RowVectorPtr input);
127127

128-
// Flush remaining data and coordinate with peer drivers. Called from
129-
// noMoreInput() and close(). Idempotent via finished_ guard.
128+
// Flush remaining data and coordinate with peer drivers.
130129
void finish();
131130

131+
// Publish buffer stats as operator runtime stats.
132+
void recordBufferStats();
133+
132134
// Partition input, serialize into flat buffer, and flush if threshold
133135
// reached.
134136
void flushBatch();

‎presto-native-execution/presto_cpp/main/operators/tests/MaterializedExchangeTest.cpp‎

Lines changed: 93 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,33 @@ std::string localShuffleReadInfo(
105105
.serialize();
106106
}
107107

108+
// ShuffleWriter that delegates all operations to a real writer but throws
109+
// from noMoreData(success=true) to simulate a writer close failure.
110+
class FailingCloseShuffleWriter : public ShuffleWriter {
111+
public:
112+
explicit FailingCloseShuffleWriter(std::shared_ptr<ShuffleWriter> delegate)
113+
: delegate_(std::move(delegate)) {}
114+
115+
void collect(int32_t partition, std::string_view key, std::string_view data)
116+
override {
117+
delegate_->collect(partition, key, data);
118+
}
119+
120+
void noMoreData(bool success) override {
121+
if (success) {
122+
VELOX_FAIL("Simulated writer close failure");
123+
}
124+
delegate_->noMoreData(success);
125+
}
126+
127+
folly::F14FastMap<std::string, int64_t> stats() const override {
128+
return delegate_->stats();
129+
}
130+
131+
private:
132+
std::shared_ptr<ShuffleWriter> delegate_;
133+
};
134+
108135
} // namespace
109136

110137
class MaterializedExchangeTest : public exec::test::OperatorTestBase {
@@ -518,6 +545,72 @@ TEST_F(MaterializedExchangeTest, replicateNullsAndAnyDisabled) {
518545
cleanupDirectory(tempDir_->getPath());
519546
}
520547

548+
// Verifies that exceptions from buffer_->noMoreData() (e.g., writer close
549+
// failure) propagate to the task instead of being silently swallowed.
550+
// Previously a broad try/catch in finish() caught these exceptions, causing the
551+
// operator to report success despite the buffer being aborted — silent data
552+
// loss.
553+
TEST_F(MaterializedExchangeTest, assertBufferCloseExceptionsArePropagated) {
554+
auto data = makeRowVector({
555+
makeFlatVector<int32_t>({1, 2, 3, 4, 5, 6}),
556+
makeFlatVector<std::string>({"a", "bb", "ccc", "dddd", "eeeee", "f"}),
557+
});
558+
559+
const int numPartitions = 4;
560+
const int numDrivers = 1;
561+
auto dataType = asRowType(data->type());
562+
563+
auto writeInfoStr = localShuffleWriteInfo(tempDir_->getPath(), numPartitions);
564+
auto writeInfo = LocalShuffleWriteInfo::deserialize(writeInfoStr);
565+
566+
constexpr uint64_t kMaxBytesPerPartition = 1 << 20;
567+
auto realWriter = std::make_shared<LocalShuffleWriter>(
568+
writeInfo.rootPath,
569+
writeInfo.queryId,
570+
writeInfo.shuffleId,
571+
writeInfo.numPartitions,
572+
kMaxBytesPerPartition,
573+
writeInfo.sortedShuffle,
574+
pool());
575+
576+
auto failingWriter =
577+
std::make_shared<FailingCloseShuffleWriter>(std::move(realWriter));
578+
579+
auto writerPool = rootPool_->addLeafChild("writerPool");
580+
581+
constexpr size_t kMaxBufferedBytes = 1 << 20;
582+
auto buffer = std::make_shared<MaterializedOutputBuffer>(
583+
numPartitions, failingWriter, writerPool, kMaxBufferedBytes);
584+
585+
std::vector<core::TypedExprPtr> keys{
586+
std::make_shared<core::FieldAccessTypedExpr>(
587+
dataType->childAt(0), dataType->nameOf(0))};
588+
589+
auto partitionFunctionSpec =
590+
std::make_shared<exec::HashPartitionFunctionSpec>(
591+
dataType, std::vector<column_index_t>{0});
592+
593+
auto valuesNode = exec::test::PlanBuilder().values({data}, true).planNode();
594+
auto exchangeWriteNode = std::make_shared<MaterializedOutputNode>(
595+
"exchangeWrite",
596+
keys,
597+
numPartitions,
598+
dataType,
599+
partitionFunctionSpec,
600+
false,
601+
valuesNode,
602+
buffer);
603+
604+
auto taskId = makeTaskId("write", 0);
605+
auto task = makeTask(taskId, exchangeWriteNode, 0);
606+
task->start(numDrivers);
607+
608+
EXPECT_TRUE(exec::test::waitForTaskFailure(task.get(), 10'000'000))
609+
<< "Task should fail when buffer noMoreData() throws, not succeed silently";
610+
611+
cleanupDirectory(tempDir_->getPath());
612+
}
613+
521614
} // namespace facebook::presto::operators::test
522615

523616
int main(int argc, char** argv) {

‎presto-native-execution/presto_cpp/main/types/PrestoToVeloxQueryPlan.cpp‎

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -39,6 +39,7 @@
3939
#include "presto_cpp/main/operators/ShuffleWrite.h"
4040
#include "presto_cpp/main/properties/session/SessionProperties.h"
4141
#include "presto_cpp/main/types/TypeParser.h"
42+
#include "velox/exec/MemoryReclaimer.h"
4243
#include "velox/exec/TraceUtil.h"
4344
// RPC plan nodes for single-operator async RPC execution
4445
#include <folly/json.h>
@@ -2650,8 +2651,11 @@ core::PlanFragment VeloxBatchQueryPlanConverter::toVeloxQueryPlan(
26502651
auto shuffleWriterPool = useSystemMemory
26512652
? velox::memory::memoryManager()->addLeafPool(
26522653
fmt::format("_sys.exchange_writer.{}", taskId))
2654+
// Add noop memory reclaimer to participate in arbitration protocol.
26532655
: queryCtx_->pool()->addLeafChild(
2654-
fmt::format("exchange_writer.{}", taskId));
2656+
fmt::format("exchange_writer.{}", taskId),
2657+
true,
2658+
velox::exec::MemoryReclaimer::create());
26552659
auto shuffleWriter = shuffleFactory->createWriter(
26562660
*serializedShuffleWriteInfo_, shuffleWriterPool.get());
26572661

0 commit comments

Comments
 (0)