Wilson parallel upload (#1875)

domwst · web-flow · commit 4a5e2ba17b29 · 2024-02-20T15:02:14.000+01:00
diff --git a/ydb/core/driver_lib/run/kikimr_services_initializers.cpp b/ydb/core/driver_lib/run/kikimr_services_initializers.cpp
@@ -870,6 +870,7 @@ void TBasicServicesInitializer::InitializeServices(NActors::TActorSystemSetup* s
                     GET_FIELD_FROM_CONFIG(MaxSpansInBatch)
                     GET_FIELD_FROM_CONFIG(MaxBytesInBatch)
                     GET_FIELD_FROM_CONFIG(SpanExportTimeoutSeconds)
+                    GET_FIELD_FROM_CONFIG(MaxExportRequestsInflight)
 
 #undef GET_FIELD_FROM_CONFIG
 
diff --git a/ydb/core/protos/config.proto b/ydb/core/protos/config.proto
@@ -1577,6 +1577,8 @@ message TTracingConfig {
         // time after which generated span will be discarded and will
         // not be sent to the collector
         optional uint32 SpanExportTimeoutSeconds = 5;
+        // maximum batch export requests being run simultaneously
+        optional uint64 MaxExportRequestsInflight = 6;
     }
 
     reserved 1 to 5;
diff --git a/ydb/library/actors/wilson/wilson_uploader.cpp b/ydb/library/actors/wilson/wilson_uploader.cpp
@@ -84,6 +84,13 @@ namespace NWilson {
             }
         };
 
+        struct TExportRequestData : TIntrusiveListItem<TExportRequestData> {
+            std::unique_ptr<grpc::ClientContext> Context;
+            std::unique_ptr<grpc::ClientAsyncResponseReader<NServiceProto::ExportTraceServiceResponse>> Reader;
+            grpc::Status Status;
+            NServiceProto::ExportTraceServiceResponse Response;
+        };
+
         class TWilsonUploader
             : public TActorBootstrapped<TWilsonUploader>
         {
@@ -95,6 +102,7 @@ namespace NWilson {
             ui64 MaxBytesInBatch;
             TDuration MaxBatchAccumulation = TDuration::Seconds(1);
             TDuration MaxSpanTimeInQueue;
+            ui64 MaxExportInflight;
 
             bool WakeupScheduled = false;
 
@@ -106,10 +114,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;
 
             TBatch CurrentBatch;
             std::queue<TBatch::TData> BatchQueue;
@@ -119,13 +123,17 @@ namespace NWilson {
             bool BatchCompletionScheduled = false;
             TMonotonic NextBatchCompletion;
 
+            TIntrusiveListWithAutoDelete<TExportRequestData, TDelete> ExportRequests;
+            size_t ExportRequestsCount = 0;
+
         public:
             TWilsonUploader(WilsonUploaderParams params)
                 : MaxSpansPerSecond(params.MaxSpansPerSecond)
                 , MaxSpansInBatch(params.MaxSpansInBatch)
                 , MaxBytesInBatch(params.MaxBytesInBatch)
                 , MaxBatchAccumulation(params.MaxBatchAccumulation)
                 , MaxSpanTimeInQueue(TDuration::Seconds(params.SpanExportTimeoutSeconds))
+                , MaxExportInflight(params.MaxExportRequestsInflight)
                 , CollectorUrl(std::move(params.CollectorUrl))
                 , ServiceName(std::move(params.ServiceName))
                 , GrpcSigner(std::move(params.GrpcSigner))
@@ -149,6 +157,10 @@ namespace NWilson {
                     ALOG_WARN(WILSON_SERVICE_ID, "max_spans_in_batch shold be greater than 0, changing to 1");
                     MaxSpansInBatch = 1;
                 }
+                if (MaxExportInflight == 0) {
+                    ALOG_WARN(WILSON_SERVICE_ID, "max_span_export_inflight should be greater than 0, changing to 1");
+                    MaxExportInflight = 1;
+                }
 
                 TStringBuf scheme;
                 TStringBuf host;
@@ -243,7 +255,7 @@ namespace NWilson {
                         "dropped " << numSpansDropped << " span(s) due to expiration");
                 }
 
-                if (Context || BatchQueue.empty()) {
+                if (ExportRequestsCount >= MaxExportInflight || BatchQueue.empty()) {
                     return;
                 } else if (now < NextSendTimestamp) {
                     ScheduleWakeup(NextSendTimestamp);
@@ -267,29 +279,42 @@ namespace NWilson {
                 SpansSizeBytes -= batch.SizeBytes;
 
                 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(batch.Request), &CQ);
-                Reader->Finish(&Response, &Status, nullptr);
+                auto reader = Stub->AsyncExport(context.get(), std::move(batch.Request), &CQ);
+                auto uploadData =  std::unique_ptr<TExportRequestData>(new TExportRequestData {
+                    .Context = std::move(context),
+                    .Reader = std::move(reader),
+                });
+                uploadData->Reader->Finish(&uploadData->Response, &uploadData->Status, uploadData.get());
+                ALOG_TRACE(WILSON_SERVICE_ID, "started export request " << (void*)uploadData.get());
+                ExportRequests.PushBack(uploadData.release());
+                ++ExportRequestsCount;
             }
 
-            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 (ExportRequests.Empty()) {
+                    return;
+                }
+                void* tag;
+                bool ok;
+                while (CQ.AsyncNext(&tag, &ok, std::chrono::system_clock::now()) == grpc::CompletionQueue::GOT_EVENT) {
+                    auto node = std::unique_ptr<TExportRequestData>(static_cast<TExportRequestData*>(tag));
+                    ALOG_TRACE(WILSON_SERVICE_ID, "finished export request " << (void*)node.get());
+                    if (!node->Status.ok()) {
+                        ALOG_ERROR(WILSON_SERVICE_ID,
+                            "failed to commit traces: " << node->Status.error_message());
                     }
+                    
+                    --ExportRequestsCount;
+                    node->Unlink();
+                }
+
+                if (!ExportRequests.Empty()) {
+                    ScheduleWakeup(TDuration::MilliSeconds(100));
                 }
             }
 
@@ -322,7 +347,7 @@ namespace NWilson {
             }
 
             void TryMakeProgress() {
-                CheckIfDone();
+                ReapCompletedRequests();
                 TryToSend();
             }
 
diff --git a/ydb/library/actors/wilson/wilson_uploader.h b/ydb/library/actors/wilson/wilson_uploader.h
@@ -36,6 +36,7 @@ namespace NWilson {
         ui64 MaxBatchAccumulationMilliseconds = 1'000;
         TDuration MaxBatchAccumulation = TDuration::Seconds(1);
         ui32 SpanExportTimeoutSeconds = 60 * 60 * 24 * 365;
+        ui64 MaxExportRequestsInflight = 1;
 
         NActors::IActor* CreateUploader() &&;
     };
diff --git a/ydb/tools/cfg/static.py b/ydb/tools/cfg/static.py
@@ -1208,7 +1208,11 @@ def get_uploader(uploader):
 
             span_export_timeout_seconds = uploader.get("span_export_timeout_seconds")
             if span_export_timeout_seconds is not None:
-                uploader_pb.SpanTtlSeconds = span_export_timeout_seconds
+                uploader_pb.SpanExportTimeoutSeconds = span_export_timeout_seconds
+
+            max_export_requests_inflight = uploader.get("max_export_requests_inflight")
+            if max_export_requests_inflight is not None:
+                uploader_pb.MaxExportRequestsInflight = max_export_requests_inflight
 
             return uploader_pb
 
diff --git a/ydb/tools/cfg/validation.py b/ydb/tools/cfg/validation.py
@@ -176,9 +176,11 @@
                 max_spans_per_second=dict(type="integer", minimum=1),
                 max_spans_in_batch=dict(type="integer", minimum=1),
                 max_bytes_in_batch=dict(type="integer"),
-                max_bytes_accumulation_milliseconds=dict(type="integer"),
+                max_batch_accumulation_milliseconds=dict(type="integer"),
                 span_export_timeout_seconds=dict(type="integer", minimum=1),
+                max_export_requests_inflight=dict(type="integer", minimum=1),
             ),
+            additionalProperties=False,
         ),
         sampling=dict(
             type="array",

Original file line number	Diff line number	Diff line change
`@@ -1577,6 +1577,8 @@ message TTracingConfig {`
`1577`	`1577`	`// time after which generated span will be discarded and will`
`1578`	`1578`	`// not be sent to the collector`
`1579`	`1579`	`optional uint32 SpanExportTimeoutSeconds = 5;`
	`1580`	`+ // maximum batch export requests being run simultaneously`
	`1581`	`+ optional uint64 MaxExportRequestsInflight = 6;`
`1580`	`1582`	`}`
`1581`	`1583`
`1582`	`1584`	`reserved 1 to 5;`