Implemented parallel upload in wilson

ydb-platform · Feb 13, 2024 · 32405e5 · 32405e5
1 parent 546c9c3
commit 32405e5
Showing 1 changed file with 39 additions and 24 deletions.
diff --git a/ydb/library/actors/wilson/wilson_uploader.cpp b/ydb/library/actors/wilson/wilson_uploader.cpp
@@ -31,10 +31,6 @@ namespace NWilson {
             grpc::CompletionQueue CQ;
 
             std::unique_ptr<IGrpcSigner> GrpcSigner;
-            std::unique_ptr<grpc::ClientContext> Context;
-            std::unique_ptr<grpc::ClientAsyncResponseReader<NServiceProto::ExportTraceServiceResponse>> Reader;
-            NServiceProto::ExportTraceServiceResponse Response;
-            grpc::Status Status;
 
             struct TSpanQueueItem {
                 TMonotonic ExpirationTimestamp;
@@ -51,6 +47,15 @@ namespace NWilson {
 
             bool WakeupScheduled = false;
 
+            struct TUploadData : TIntrusiveListItem<TUploadData> {
+                std::unique_ptr<grpc::ClientContext> Context;
+                std::unique_ptr<grpc::ClientAsyncResponseReader<NServiceProto::ExportTraceServiceResponse>> Reader;
+                grpc::Status Status;
+                NServiceProto::ExportTraceServiceResponse Response;
+            };
+
+            TIntrusiveListWithAutoDelete<TUploadData, TDelete> UploadRequests;
+
         public:
             TWilsonUploader(WilsonUploaderParams params)
                 : CollectorUrl(std::move(params.CollectorUrl))
@@ -131,7 +136,7 @@ namespace NWilson {
                         "dropped " << numSpansDropped << " span(s) due to expiration");
                 }
 
-                if (Context || Spans.empty()) {
+                if (Spans.empty()) {
                     return;
                 } else if (now < NextSendTimestamp) {
                     ScheduleWakeup(NextSendTimestamp);
@@ -162,29 +167,39 @@ namespace NWilson {
                 }
 
                 ScheduleWakeup(NextSendTimestamp);
-                Context = std::make_unique<grpc::ClientContext>();
+
+                auto context = std::make_unique<grpc::ClientContext>();
                 if (GrpcSigner) {
-                    GrpcSigner->SignClientContext(*Context);
+                    GrpcSigner->SignClientContext(*context);
                 }
-                Reader = Stub->AsyncExport(Context.get(), std::move(request), &CQ);
-                Reader->Finish(&Response, &Status, nullptr);
+                auto reader = Stub->AsyncExport(context.get(), std::move(request), &CQ);
+                auto uploadData = new TUploadData {
+                    .Context = std::move(context),
+                    .Reader = std::move(reader),
+                };
+                uploadData->Reader->Finish(&uploadData->Response, &uploadData->Status, uploadData);
+                UploadRequests.PushBack(uploadData);
             }
 
-            void CheckIfDone() {
-                if (Context) {
-                    void *tag;
-                    bool ok;
-                    if (CQ.AsyncNext(&tag, &ok, std::chrono::system_clock::now()) == grpc::CompletionQueue::GOT_EVENT) {
-                        if (!Status.ok()) {
-                            ALOG_ERROR(WILSON_SERVICE_ID,
-                                "failed to commit traces: " << Status.error_message());
-                        }
-
-                        Reader.reset();
-                        Context.reset();
-                    } else {
-                        ScheduleWakeup(TDuration::MilliSeconds(100));
+            void ReapCompletedRequests() {
+                if (UploadRequests.Empty()) {
+                    return;
+                }
+                void* tag;
+                bool ok;
+                while (CQ.AsyncNext(&tag, &ok, std::chrono::system_clock::now()) == grpc::CompletionQueue::GOT_EVENT) {
+                    auto node = static_cast<TUploadData*>(tag);
+                    if (!node->Status.ok()) {
+                        ALOG_ERROR(WILSON_SERVICE_ID,
+                            "failed to commit traces: " << node->Status.error_message());
                     }
+
+                    node->Unlink();
+                    delete node;
+                }
+
+                if (!UploadRequests.Empty()) {
+                    ScheduleWakeup(TDuration::MilliSeconds(100));
                 }
             }
 
@@ -204,7 +219,7 @@ namespace NWilson {
             }
 
             void TryMakeProgress() {
-                CheckIfDone();
+                ReapCompletedRequests();
                 TryToSend();
             }