Skip to content

Commit 5f386ef

Browse files
dante159753dante159753
committed
Fix Posix AIO timeout buffer lifetime
1 parent e3dc7e0 commit 5f386ef

4 files changed

Lines changed: 121 additions & 13 deletions

File tree

‎ucm/store/posix/cc/aio_impl.cc‎

Lines changed: 19 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -139,8 +139,12 @@ static inline void AioSetEventFd(struct iocb* iocb, int32_t eventfd)
139139
iocb->aio_flags |= (1 << 0) /* IOCB_FLAG_RESFD */;
140140
iocb->aio_resfd = eventfd;
141141
}
142+
struct AioCompletionData {
143+
AioImpl::Callback callback;
144+
std::shared_ptr<void> bufferOwner;
145+
};
142146
static inline void ReleaseToAioCompletion(std::unique_ptr<struct iocb>& cb,
143-
std::unique_ptr<AioImpl::Callback>& data)
147+
std::unique_ptr<AioCompletionData>& data)
144148
{
145149
data.release();
146150
cb.release();
@@ -208,7 +212,9 @@ Status AioImpl::Setup(size_t timeoutMs)
208212
Status AioImpl::ReadAsync(Io&& io)
209213
{
210214
auto cb = std::make_unique<struct iocb>();
211-
auto data = std::make_unique<Callback>(std::move(io.callback));
215+
auto data = std::make_unique<AioCompletionData>();
216+
data->callback = std::move(io.callback);
217+
data->bufferOwner = std::move(io.bufferOwner);
212218
AioPrepareRead(cb.get(), io.fd, io.buffer, io.length, io.offset);
213219
cb->aio_data = (uintptr_t)(void*)data.get();
214220
Track(io.tag, cb.get());
@@ -225,7 +231,9 @@ Status AioImpl::ReadAsync(Io&& io)
225231
Status AioImpl::WriteAsync(Io&& io)
226232
{
227233
auto cb = std::make_unique<struct iocb>();
228-
auto data = std::make_unique<Callback>(std::move(io.callback));
234+
auto data = std::make_unique<AioCompletionData>();
235+
data->callback = std::move(io.callback);
236+
data->bufferOwner = std::move(io.bufferOwner);
229237
AioPrepareWrite(cb.get(), io.fd, io.buffer, io.length, io.offset);
230238
cb->aio_data = (uintptr_t)(void*)data.get();
231239
Track(io.tag, cb.get());
@@ -276,7 +284,7 @@ void AioImpl::HarvestCompletions(std::vector<io_event>& events)
276284
auto num = AioGetEvents(ctx_, 1, batchSize, events.data(), nullptr);
277285
for (auto i = 0; i < num; i++) {
278286
auto* iocbPtr = reinterpret_cast<struct iocb*>(events[i].obj);
279-
auto cb = (Callback*)(void*)events[i].data;
287+
auto data = (AioCompletionData*)(void*)events[i].data;
280288
Result res;
281289
if (events[i].res >= 0) {
282290
res.nBytes = events[i].res;
@@ -285,9 +293,9 @@ void AioImpl::HarvestCompletions(std::vector<io_event>& events)
285293
res.nBytes = -1;
286294
res.error = -static_cast<int>(events[i].res);
287295
}
288-
if (cb) {
289-
(*cb)(res);
290-
delete cb;
296+
if (data) {
297+
data->callback(res);
298+
delete data;
291299
}
292300
if (iocbPtr) {
293301
Untrack(iocbPtr);
@@ -347,10 +355,10 @@ void AioImpl::CancelTask(uint64_t tag)
347355
if (vec.empty()) { iocbTable_.erase(it); }
348356
}
349357
for (auto* cbp : cancelled) {
350-
auto* cb = reinterpret_cast<Callback*>(cbp->aio_data);
351-
if (cb) {
352-
(*cb)(Result{-1, ECANCELED});
353-
delete cb;
358+
auto* data = reinterpret_cast<AioCompletionData*>(cbp->aio_data);
359+
if (data) {
360+
data->callback(Result{-1, ECANCELED});
361+
delete data;
354362
}
355363
delete cbp;
356364
}

‎ucm/store/posix/cc/aio_impl.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
#include <cstdint>
2929
#include <functional>
3030
#include <linux/aio_abi.h>
31+
#include <memory>
3132
#include <mutex>
3233
#include <thread>
3334
#include <unordered_map>
@@ -58,6 +59,7 @@ class AioImpl {
5859
uint64_t offset;
5960
uint32_t length;
6061
void* buffer;
62+
std::shared_ptr<void> bufferOwner;
6163
Callback callback;
6264
uint64_t tag{0};
6365
};

‎ucm/store/posix/cc/io_engine_aio.h‎

Lines changed: 31 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,10 @@
2525
#define UNIFIEDCACHE_POSIX_STORE_CC_IO_ENGINE_AIO_H
2626

2727
#include <cerrno>
28+
#include <cstdlib>
29+
#include <cstring>
2830
#include <limits>
31+
#include <memory>
2932
#include <mutex>
3033
#include <tuple>
3134
#include <unordered_map>
@@ -95,6 +98,12 @@ class IoEngineAio : public Detail::TaskWrapper<TransTask, Detail::TaskHandle> {
9598
static UC::Metrics::CachedMetric ioErrors{"posix_io_errors_total"};
9699
UC::Metrics::UpdateStats(ioErrors, 1.0);
97100
}
101+
std::shared_ptr<void> MakeOwnedIoBuffer() const
102+
{
103+
void* buffer = nullptr;
104+
if (posix_memalign(&buffer, 4096, shardSize_) != 0) { return {}; }
105+
return std::shared_ptr<void>(buffer, [](void* ptr) { std::free(ptr); });
106+
}
98107
void CommitBlock(Detail::BlockId id, bool success)
99108
{
100109
blockOperator_.Submit(BlockOperator::CommitTask{std::move(id), success});
@@ -142,13 +151,33 @@ class IoEngineAio : public Detail::TaskWrapper<TransTask, Detail::TaskHandle> {
142151
w->Done();
143152
return;
144153
}
154+
auto* userBuffer = shard.addrs.front();
155+
auto ioBuffer = userBuffer;
156+
std::shared_ptr<void> ioBufferOwner;
157+
if (timeoutMs_ > 0) {
158+
ioBufferOwner = MakeOwnedIoBuffer();
159+
if (!ioBufferOwner) {
160+
UC_ERROR("Failed to allocate owned aio buffer for task({}).", tid);
161+
handleFailure(result.fd);
162+
return;
163+
}
164+
ioBuffer = ioBufferOwner.get();
165+
if constexpr (dump) { std::memcpy(ioBuffer, userBuffer, shardSize_); }
166+
}
145167
AioImpl::Io io;
146168
io.fd = result.fd;
147169
io.offset = shard.index * shardSize_;
148170
io.length = shardSize_;
149-
io.buffer = shard.addrs.front();
171+
io.buffer = ioBuffer;
172+
io.bufferOwner = ioBufferOwner;
150173
io.tag = tid;
151-
io.callback = [this, tid, w, fd = result.fd, last, id](AioImpl::Result ioResult) {
174+
io.callback = [this, tid, w, fd = result.fd, last, id, ioBuffer, userBuffer,
175+
copySize = shardSize_](AioImpl::Result ioResult) {
176+
if constexpr (!dump) {
177+
if (ioResult.error == 0 && !failureSet_.Contains(tid)) {
178+
std::memcpy(userBuffer, ioBuffer, copySize);
179+
}
180+
}
152181
OnIoCallback<dump>(tid, w, fd, last, id, ioResult);
153182
};
154183
auto status = dump ? aio_.WriteAsync(std::move(io)) : aio_.ReadAsync(std::move(io));

‎ucm/store/test/case/posix/posix_store_test.cc‎

Lines changed: 69 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -316,6 +316,75 @@ TEST_F(UCPosixStoreTest, AioWaitTimesOutWhenCompletionIsLost)
316316
ASSERT_LT(std::chrono::duration_cast<std::chrono::milliseconds>(elapsed).count(), 1000);
317317
}
318318

319+
TEST_F(UCPosixStoreTest, AioDumpSubmitsOwnedBufferWhenTimeoutCanReturn)
320+
{
321+
using namespace UC::PosixStore;
322+
PosixStore store;
323+
ASSERT_EQ(store.Setup(MakeAioConfig(Path())), UC::Status::OK());
324+
ScopedAioHooks hooks;
325+
std::atomic<uintptr_t> submittedBuffer{0};
326+
TestHooks::SetAioSubmitHook([&submittedBuffer](aio_context_t, int64_t nr, iocb** ios) {
327+
submittedBuffer.store(ios[0]->aio_buf, std::memory_order_relaxed);
328+
return static_cast<int32_t>(nr);
329+
});
330+
TestHooks::SetAioCancelHook([](aio_context_t, struct iocb*, io_event*) { return 0; });
331+
auto buffer = MakeAlignedBuffer(8);
332+
ASSERT_NE(buffer.get(), nullptr);
333+
auto block = UC::Test::Detail::TypesHelper::MakeBlockIdRandomly();
334+
auto handle = store.Dump(MakeDumpDesc("AioDumpOwnedBuffer", block, buffer.get()));
335+
ASSERT_TRUE(handle.HasValue());
336+
337+
ASSERT_EQ(store.Wait(handle.Value()), UC::Status::Timeout());
338+
ASSERT_NE(submittedBuffer.load(std::memory_order_relaxed), 0U);
339+
ASSERT_NE(submittedBuffer.load(std::memory_order_relaxed),
340+
reinterpret_cast<uintptr_t>(buffer.get()));
341+
}
342+
343+
TEST_F(UCPosixStoreTest, AioLoadSubmitsOwnedBufferWhenTimeoutCanReturn)
344+
{
345+
using namespace UC::PosixStore;
346+
auto block = UC::Test::Detail::TypesHelper::MakeBlockIdRandomly();
347+
auto dumpBuffer = MakeAlignedBuffer(9);
348+
ASSERT_NE(dumpBuffer.get(), nullptr);
349+
{
350+
PosixStore store;
351+
ASSERT_EQ(store.Setup(MakeAioConfig(Path(), 5000)), UC::Status::OK());
352+
auto handle = store.Dump(MakeDumpDesc("AioLoadOwnedBufferPrepare", block, dumpBuffer.get()));
353+
ASSERT_TRUE(handle.HasValue());
354+
ASSERT_EQ(store.Wait(handle.Value()), UC::Status::OK());
355+
bool found = false;
356+
for (size_t i = 0; i < 20; ++i) {
357+
auto founds = store.Lookup(&block, 1);
358+
ASSERT_TRUE(founds.HasValue());
359+
if (founds.Value()[0]) {
360+
found = true;
361+
break;
362+
}
363+
std::this_thread::sleep_for(std::chrono::milliseconds(50));
364+
}
365+
ASSERT_TRUE(found);
366+
}
367+
368+
PosixStore store;
369+
ASSERT_EQ(store.Setup(MakeAioConfig(Path())), UC::Status::OK());
370+
ScopedAioHooks hooks;
371+
std::atomic<uintptr_t> submittedBuffer{0};
372+
TestHooks::SetAioSubmitHook([&submittedBuffer](aio_context_t, int64_t nr, iocb** ios) {
373+
submittedBuffer.store(ios[0]->aio_buf, std::memory_order_relaxed);
374+
return static_cast<int32_t>(nr);
375+
});
376+
TestHooks::SetAioCancelHook([](aio_context_t, struct iocb*, io_event*) { return 0; });
377+
auto loadBuffer = MakeAlignedBuffer(0);
378+
ASSERT_NE(loadBuffer.get(), nullptr);
379+
auto handle = store.Load(MakeDumpDesc("AioLoadOwnedBuffer", block, loadBuffer.get()));
380+
ASSERT_TRUE(handle.HasValue());
381+
382+
ASSERT_EQ(store.Wait(handle.Value()), UC::Status::Timeout());
383+
ASSERT_NE(submittedBuffer.load(std::memory_order_relaxed), 0U);
384+
ASSERT_NE(submittedBuffer.load(std::memory_order_relaxed),
385+
reinterpret_cast<uintptr_t>(loadBuffer.get()));
386+
}
387+
319388
TEST_F(UCPosixStoreTest, AioCheckFinishesLostCompletionAfterDeadline)
320389
{
321390
using namespace UC::PosixStore;

0 commit comments

Comments
 (0)