update: check pipeline queue before scrape, try again after 1 second …

…if is not valid to push (alibaba#1757)
quzard · Oct 8, 2024 · 6d9d5f2 · 6d9d5f2
1 parent 6526c34
commit 6d9d5f2
Show file tree

Hide file tree

Showing 13 changed files with 119 additions and 37 deletions.
diff --git a/core/common/timer/Timer.cpp b/core/common/timer/Timer.cpp
@@ -66,12 +66,15 @@ void Timer::Run() {
                     mCV.wait_for(threadLock, timeout);
                     break;
                 } else {
+                    auto e = std::move(const_cast<unique_ptr<TimerEvent>&>(mQueue.top()));
+                    mQueue.pop();
+                    queueLock.unlock();
                     if (!e->IsValid()) {
                         LOG_INFO(sLogger, ("invalid timer event", "task is cancelled"));
                     } else {
                         e->Execute();
                     }
-                    mQueue.pop();
+                    queueLock.lock();
                 }
             }
         }

diff --git a/core/common/timer/Timer.h b/core/common/timer/Timer.h
@@ -52,6 +52,7 @@ class Timer {
 
 #ifdef APSARA_UNIT_TEST_MAIN
     friend class TimerUnittest;
+    friend class ScrapeSchedulerUnittest;
 #endif
 };
 

diff --git a/core/plugin/input/InputPrometheus.h b/core/plugin/input/InputPrometheus.h
@@ -17,7 +17,7 @@ class InputPrometheus : public Input {
     bool Init(const Json::Value& config, Json::Value& optionalGoPipeline) override;
     bool Start() override;
     bool Stop(bool isPipelineRemoving) override;
-    bool SupportAck() const override { return false; }
+    bool SupportAck() const override { return true; }
 
 private:
     bool CreateInnerProcessors(const Json::Value& inputConfig);

diff --git a/core/prometheus/async/PromFuture.cpp b/core/prometheus/async/PromFuture.cpp
@@ -1,28 +1,38 @@
 #include "prometheus/async/PromFuture.h"
 
 #include "common/Lock.h"
+#include "common/http/HttpResponse.h"
 
 namespace logtail {
 
-void PromFuture::Process(const HttpResponse& response, uint64_t timestampMilliSec) {
+template <typename... Args>
+bool PromFuture<Args...>::Process(Args... args) {
     WriteLock lock(mStateRWLock);
     if (mState == PromFutureState::New) {
         for (auto& callback : mDoneCallbacks) {
-            callback(response, timestampMilliSec);
+            if (!callback(std::forward<Args>(args)...)) {
+                mState = PromFutureState::Done;
+                return false;
+            }
         }
         mState = PromFutureState::Done;
-    } else {
-        return;
     }
+
+    return true;
 }
 
-void PromFuture::AddDoneCallback(std::function<void(const HttpResponse&, uint64_t timestampMilliSec)>&& callback) {
+template <typename... Args>
+void PromFuture<Args...>::AddDoneCallback(CallbackSignature&& callback) {
     mDoneCallbacks.emplace_back(std::move(callback));
 }
 
-void PromFuture::Cancel() {
+template <typename... Args>
+void PromFuture<Args...>::Cancel() {
     WriteLock lock(mStateRWLock);
     mState = PromFutureState::Done;
 }
 
+template class PromFuture<const HttpResponse&, uint64_t>;
+template class PromFuture<>;
+
 } // namespace logtail
diff --git a/core/prometheus/async/PromFuture.h b/core/prometheus/async/PromFuture.h
@@ -1,26 +1,29 @@
 #pragma once
 
+#include <functional>
+
 #include "common/Lock.h"
-#include "common/http/HttpResponse.h"
 
 namespace logtail {
 
 enum class PromFutureState { New, Processing, Done };
 
+template <typename... Args>
 class PromFuture {
 public:
+    using CallbackSignature = std::function<bool(Args...)>;
     // Process should support oneshot and streaming mode.
-    void Process(const HttpResponse&, uint64_t timestampMilliSec);
+    bool Process(Args...);
 
-    void AddDoneCallback(std::function<void(const HttpResponse&, uint64_t timestampMilliSec)>&& callback);
+    void AddDoneCallback(CallbackSignature&&);
 
     void Cancel();
 
 protected:
     PromFutureState mState = {PromFutureState::New};
     ReadWriteLock mStateRWLock;
 
-    std::vector<std::function<void(const HttpResponse&, uint64_t timestampMilliSec)>> mDoneCallbacks;
+    std::vector<CallbackSignature> mDoneCallbacks;
 
 #ifdef APSARA_UNIT_TEST_MAIN
     friend class ScrapeSchedulerUnittest;

diff --git a/core/prometheus/async/PromHttpRequest.cpp b/core/prometheus/async/PromHttpRequest.cpp
@@ -18,17 +18,24 @@ PromHttpRequest::PromHttpRequest(const std::string& method,
                                  const std::string& body,
                                  uint32_t timeout,
                                  uint32_t maxTryCnt,
-                                 std::shared_ptr<PromFuture> future)
+                                 std::shared_ptr<PromFuture<const HttpResponse&, uint64_t>> future,
+                                 std::shared_ptr<PromFuture<>> isContextValidFuture)
     : AsynHttpRequest(method, httpsFlag, host, port, url, query, header, body, timeout, maxTryCnt),
-      mFuture(std::move(future)) {
+      mFuture(std::move(future)),
+      mIsContextValidFuture(std::move(isContextValidFuture)) {
 }
 
 void PromHttpRequest::OnSendDone(const HttpResponse& response) {
-    mFuture->Process(response,
-                     std::chrono::duration_cast<std::chrono::milliseconds>(mLastSendTime.time_since_epoch()).count());
+    if (mFuture != nullptr) {
+        mFuture->Process(
+            response, std::chrono::duration_cast<std::chrono::milliseconds>(mLastSendTime.time_since_epoch()).count());
+    }
 }
 
 [[nodiscard]] bool PromHttpRequest::IsContextValid() const {
+    if (mIsContextValidFuture != nullptr) {
+        return mIsContextValidFuture->Process();
+    }
     return true;
 }
 

diff --git a/core/prometheus/async/PromHttpRequest.h b/core/prometheus/async/PromHttpRequest.h
@@ -20,7 +20,8 @@ class PromHttpRequest : public AsynHttpRequest {
                     const std::string& body,
                     uint32_t timeout,
                     uint32_t maxTryCnt,
-                    std::shared_ptr<PromFuture> future);
+                    std::shared_ptr<PromFuture<const HttpResponse&, uint64_t>> future,
+                    std::shared_ptr<PromFuture<>> isContextValidFuture = nullptr);
     PromHttpRequest(const PromHttpRequest&) = default;
     ~PromHttpRequest() override = default;
 
@@ -30,7 +31,8 @@ class PromHttpRequest : public AsynHttpRequest {
 private:
     void SetNextExecTime(std::chrono::steady_clock::time_point execTime);
 
-    std::shared_ptr<PromFuture> mFuture;
+    std::shared_ptr<PromFuture<const HttpResponse&, uint64_t>> mFuture;
+    std::shared_ptr<PromFuture<>> mIsContextValidFuture;
 };
 
 } // namespace logtail
diff --git a/core/prometheus/schedulers/BaseScheduler.cpp b/core/prometheus/schedulers/BaseScheduler.cpp
@@ -3,14 +3,20 @@
 namespace logtail {
 void BaseScheduler::ExecDone() {
     mExecCount++;
+    mLatestExecTime = mFirstExecTime + std::chrono::seconds(mExecCount * mInterval);
 }
 
 std::chrono::steady_clock::time_point BaseScheduler::GetNextExecTime() {
-    return mFirstExecTime + std::chrono::seconds(mExecCount * mInterval);
+    return mLatestExecTime;
 }
 
 void BaseScheduler::SetFirstExecTime(std::chrono::steady_clock::time_point firstExecTime) {
     mFirstExecTime = firstExecTime;
+    mLatestExecTime = mFirstExecTime;
+}
+
+void BaseScheduler::DelayExecTime(uint64_t delaySeconds) {
+    mLatestExecTime = mLatestExecTime + std::chrono::seconds(delaySeconds);
 }
 
 void BaseScheduler::Cancel() {

diff --git a/core/prometheus/schedulers/BaseScheduler.h b/core/prometheus/schedulers/BaseScheduler.h
@@ -3,6 +3,7 @@
 
 #include <memory>
 
+#include "common/http/HttpResponse.h"
 #include "prometheus/async/PromFuture.h"
 
 namespace logtail {
@@ -18,18 +19,21 @@ class BaseScheduler {
     std::chrono::steady_clock::time_point GetNextExecTime();
 
     void SetFirstExecTime(std::chrono::steady_clock::time_point firstExecTime);
+    void DelayExecTime(uint64_t delaySeconds);
 
     virtual void Cancel();
 
 protected:
     bool IsCancelled();
 
     std::chrono::steady_clock::time_point mFirstExecTime;
+    std::chrono::steady_clock::time_point mLatestExecTime;
     int64_t mExecCount = 0;
     int64_t mInterval = 0;
 
     ReadWriteLock mLock;
     bool mValidState = true;
-    std::shared_ptr<PromFuture> mFuture;
+    std::shared_ptr<PromFuture<const HttpResponse&, uint64_t>> mFuture;
+    std::shared_ptr<PromFuture<>> mIsContextValidFuture;
 };
 } // namespace logtail
diff --git a/core/prometheus/schedulers/ScrapeScheduler.cpp b/core/prometheus/schedulers/ScrapeScheduler.cpp
@@ -31,6 +31,7 @@
 #include "pipeline/queue/ProcessQueueManager.h"
 #include "pipeline/queue/QueueKey.h"
 #include "prometheus/Constants.h"
+#include "prometheus/async/PromFuture.h"
 #include "prometheus/async/PromHttpRequest.h"
 
 using namespace std;
@@ -99,40 +100,60 @@ void ScrapeScheduler::PushEventGroup(PipelineEventGroup&& eGroup) {
     auto item = make_unique<ProcessQueueItem>(std::move(eGroup), mInputIndex);
 #ifdef APSARA_UNIT_TEST_MAIN
     mItem.push_back(std::move(item));
+    return;
 #endif
-    ProcessQueueManager::GetInstance()->PushQueue(mQueueKey, std::move(item));
+    while (true) {
+        if (ProcessQueueManager::GetInstance()->PushQueue(mQueueKey, std::move(item)) == 0) {
+            break;
+        }
+        usleep(10 * 1000);
+    }
 }
 
 string ScrapeScheduler::GetId() const {
     return mHash;
 }
 
 void ScrapeScheduler::ScheduleNext() {
-    auto future = std::make_shared<PromFuture>();
+    auto future = std::make_shared<PromFuture<const HttpResponse&, uint64_t>>();
+    auto isContextValidFuture = std::make_shared<PromFuture<>>();
     future->AddDoneCallback([this](const HttpResponse& response, uint64_t timestampMilliSec) {
         this->OnMetricResult(response, timestampMilliSec);
         this->ExecDone();
         this->ScheduleNext();
+        return true;
+    });
+    isContextValidFuture->AddDoneCallback([this]() -> bool {
+        if (ProcessQueueManager::GetInstance()->IsValidToPush(mQueueKey)) {
+            return true;
+        } else {
+            this->DelayExecTime(1);
+            this->ScheduleNext();
+            return false;
+        }
     });
 
     if (IsCancelled()) {
         mFuture->Cancel();
+        mIsContextValidFuture->Cancel();
         return;
     }
 
     {
         WriteLock lock(mLock);
         mFuture = future;
+        mIsContextValidFuture = isContextValidFuture;
     }
 
     auto event = BuildScrapeTimerEvent(GetNextExecTime());
     mTimer->PushEvent(std::move(event));
 }
 
 void ScrapeScheduler::ScrapeOnce(std::chrono::steady_clock::time_point execTime) {
-    auto future = std::make_shared<PromFuture>();
+    auto future = std::make_shared<PromFuture<const HttpResponse&, uint64_t>>();
     future->AddDoneCallback([this](const HttpResponse& response, uint64_t timestampMilliSec) {
         this->OnMetricResult(response, timestampMilliSec);
+        return true;
     });
     mFuture = future;
     auto event = BuildScrapeTimerEvent(execTime);
@@ -153,15 +174,19 @@ std::unique_ptr<TimerEvent> ScrapeScheduler::BuildScrapeTimerEvent(std::chrono::
                                                      mScrapeConfigPtr->mScrapeTimeoutSeconds,
                                                      mScrapeConfigPtr->mScrapeIntervalSeconds
                                                          / mScrapeConfigPtr->mScrapeTimeoutSeconds,
-                                                     this->mFuture);
+                                                     this->mFuture,
+                                                     this->mIsContextValidFuture);
     auto timerEvent = std::make_unique<HttpRequestTimerEvent>(execTime, std::move(request));
     return timerEvent;
 }
 
 void ScrapeScheduler::Cancel() {
-    if (mFuture) {
+    if (mFuture != nullptr) {
         mFuture->Cancel();
     }
+    if (mIsContextValidFuture != nullptr) {
+        mIsContextValidFuture->Cancel();
+    }
     {
         WriteLock lock(mLock);
         mValidState = false;

diff --git a/core/prometheus/schedulers/TargetSubscriberScheduler.cpp b/core/prometheus/schedulers/TargetSubscriberScheduler.cpp
@@ -237,11 +237,12 @@ string TargetSubscriberScheduler::GetId() const {
 }
 
 void TargetSubscriberScheduler::ScheduleNext() {
-    auto future = std::make_shared<PromFuture>();
+    auto future = std::make_shared<PromFuture<const HttpResponse&, uint64_t>>();
     future->AddDoneCallback([this](const HttpResponse& response, uint64_t timestampMilliSec) {
         this->OnSubscription(response, timestampMilliSec);
         this->ExecDone();
         this->ScheduleNext();
+        return true;
     });
     if (IsCancelled()) {
         mFuture->Cancel();
@@ -267,9 +268,10 @@ void TargetSubscriberScheduler::Cancel() {
 }
 
 void TargetSubscriberScheduler::SubscribeOnce(std::chrono::steady_clock::time_point execTime) {
-    auto future = std::make_shared<PromFuture>();
+    auto future = std::make_shared<PromFuture<const HttpResponse&, uint64_t>>();
     future->AddDoneCallback([this](const HttpResponse& response, uint64_t timestampNanoSec) {
         this->OnSubscription(response, timestampNanoSec);
+        return true;
     });
     mFuture = future;
     auto event = BuildSubscriberTimerEvent(execTime);

diff --git a/core/unittest/prometheus/PromAsynUnittest.cpp b/core/unittest/prometheus/PromAsynUnittest.cpp
@@ -13,14 +13,15 @@ class PromAsynUnittest : public testing::Test {
 };
 
 void PromAsynUnittest::TestExecTime() {
-    auto future = std::make_shared<PromFuture>();
+    auto future = std::make_shared<PromFuture<const HttpResponse&, uint64_t>>();
     auto now = std::chrono::system_clock::now();
     bool exec = false;
     future->AddDoneCallback([&exec, now](const HttpResponse&, uint64_t timestampMilliSec) {
         APSARA_TEST_EQUAL(timestampMilliSec,
                           std::chrono::duration_cast<std::chrono::milliseconds>(now.time_since_epoch()).count());
 
         APSARA_TEST_TRUE(exec);
+        return true;
     });
     auto request = std::make_shared<PromHttpRequest>(
         "http", false, "127.0.0.1", 8080, "/", "", map<string, string>(), "", 10, 3, future);
-Original file line number
+Diff line change
@@ Expand Up / @@ -52,6 +52,7 @@ class Timer { @@
     #ifdef APSARA_UNIT_TEST_MAIN
         friend class TimerUnittest;
+        friend class ScrapeSchedulerUnittest;
     #endif
     };
@@ Expand Down @@