diff --git a/lix/libstore/filetransfer.cc b/lix/libstore/filetransfer.cc index d1268d98a..1b3e036a5 100644 --- a/lix/libstore/filetransfer.cc +++ b/lix/libstore/filetransfer.cc @@ -22,6 +22,7 @@ #include #include #include +#include #if ENABLE_DTRACE #include "trace-probes.gen.hh" @@ -553,7 +554,7 @@ struct TransferItem } }; -struct curlFileTransfer : public FileTransfer +struct CurlMulti { std::unique_ptr curlm; @@ -593,12 +594,12 @@ struct curlFileTransfer : public FileTransfer std::thread workerThread; - curlFileTransfer(unsigned int baseRetryTimeMs) + CurlMulti(unsigned int baseRetryTimeMs) : curlm(curl_multi_init()) , baseRetryTimeMs(baseRetryTimeMs) { if (curlm == nullptr) { - throw FileTransferError(Misc, {}, "could not allocate curl handle"); + throw FileTransferError(FileTransfer::Misc, {}, "could not allocate curl handle"); } static std::once_flag globalInit; @@ -614,7 +615,7 @@ struct curlFileTransfer : public FileTransfer }); } - ~curlFileTransfer() + ~CurlMulti() { try { stopWorkerThread(); @@ -779,6 +780,21 @@ struct curlFileTransfer : public FileTransfer } wakeup(); } +}; + +struct curlFileTransfer : public FileTransfer +{ + std::shared_ptr multi; + + curlFileTransfer(unsigned int baseRetryTimeMs) + : multi(std::make_shared(baseRetryTimeMs)) + { + } + + ~curlFileTransfer() + { + multi->stopWorkerThread(); + } #if ENABLE_S3 static std::tuple parseS3Uri(std::string uri) @@ -952,7 +968,7 @@ struct curlFileTransfer : public FileTransfer struct TransferStream : AsyncInputStream { - curlFileTransfer & parent; + std::shared_ptr parent; std::string uri; FileTransferOptions options; std::optional data; @@ -977,7 +993,7 @@ struct curlFileTransfer : public FileTransfer bool noBody, const Activity * context ) - : parent(parent) + : parent(parent.multi) , uri(uri) , options(options) , data(std::move(data)) @@ -987,7 +1003,7 @@ struct curlFileTransfer : public FileTransfer fileTransferSettings.tries, std::chrono::seconds(fileTransferSettings.maxConnectTimeout.get()), std::chrono::seconds(fileTransferSettings.initialConnectTimeout.get()), - std::chrono::milliseconds(parent.baseRetryTimeMs) + std::chrono::milliseconds(this->parent->baseRetryTimeMs) )) { } @@ -997,7 +1013,7 @@ struct curlFileTransfer : public FileTransfer // wake up the download thread if it's still going and have it abort try { if (transfer) { - parent.cancel(transfer); + parent->cancel(transfer); } } catch (...) { ignoreExceptionInDestructor(); @@ -1069,7 +1085,7 @@ struct curlFileTransfer : public FileTransfer std::move(pfp.fulfiller), timeout ); - parent.enqueueItem(transfer); + parent->enqueueItem(transfer); co_return TRY_AWAIT(pfp.promise); } catch (...) { co_return result::current_exception(); @@ -1152,13 +1168,13 @@ struct curlFileTransfer : public FileTransfer chunk = std::move(state->data); buffered = chunk; totalReceived += chunk.size(); - parent.unpause(transfer); + parent->unpause(transfer); } else if (state->exc) { std::rethrow_exception(state->exc); } else if (state->done) { co_return false; } else { - parent.unpause(transfer); + parent->unpause(transfer); signal = state->wait(); } } @@ -1246,8 +1262,9 @@ ref getFileTransfer() { static ref fileTransfer = makeCurlFileTransfer({}); - if (fileTransfer->state_.lock()->quit) + if (fileTransfer->multi->state_.lock()->quit) { fileTransfer = makeCurlFileTransfer({}); + } return fileTransfer; } diff --git a/tests/unit/libstore/filetransfer.cc b/tests/unit/libstore/filetransfer.cc index 7d4f227a8..92f47a1f1 100644 --- a/tests/unit/libstore/filetransfer.cc +++ b/tests/unit/libstore/filetransfer.cc @@ -470,4 +470,27 @@ TEST(FileTransfer, setupErrorsAreMetadata) ASSERT_THROW(aio.blockOn(ft->upload(fmt("http://[::1]:%d", port), "")), FileTransferError); } +TEST(FileTransfer, shutdownKillsTransfers) +{ + auto [port, srv] = serveHTTP({ + {"200 ok", "content-length: 999999999\r\n", [&](int) { return std::string(1024, 'X'); }}, + }); + AsyncIoRoot aio; + char buf; + std::optional> s; + { + auto ft = makeFileTransfer(0); + auto [_r, stream] = aio.blockOn(ft->download(fmt("http://[::1]:%d/index", port))); + ASSERT_EQ(aio.blockOn(stream->read(&buf, 1)), 1); + s = std::move(stream); + } + ASSERT_THROW( + { + while (true) { + aio.blockOn((*s)->drain()); + } + }, + FileTransferError + ); +} }