Intermediate changes

9 files changed, 36 insertions, 14 deletions

diff --git a/yt/yt/core/concurrency/nonblocking_batcher-inl.h b/yt/yt/core/concurrency/nonblocking_batcher-inl.h
index 6dcabdaaa3..b95df3cc88 100644
--- a/yt/yt/core/concurrency/nonblocking_batcher-inl.h
+++ b/yt/yt/core/concurrency/nonblocking_batcher-inl.h
@@ -79,8 +79,8 @@ TFuture<typename TNonblockingBatcher<T, TBatchLimiter>::TBatch> TNonblockingBatc
     auto guard = Guard(SpinLock_);
     auto promise = NewPromise<TBatch>();
     Promises_.push_back(promise);
-    StartTimer(guard);
     CheckReturn(guard);
+    StartTimer(guard);
     return promise.ToFuture();
 }
 
diff --git a/yt/yt/core/concurrency/unittests/nonblocking_batcher_ut.cpp b/yt/yt/core/concurrency/unittests/nonblocking_batcher_ut.cpp
index 8222a1e922..a06390dd2b 100644
--- a/yt/yt/core/concurrency/unittests/nonblocking_batcher_ut.cpp
+++ b/yt/yt/core/concurrency/unittests/nonblocking_batcher_ut.cpp
@@ -313,6 +313,25 @@ TEST(TNonblockingBatcherTest, AllowEmptyBatches)
     ASSERT_EQ(e2.Get().ValueOrThrow(), std::vector<int>({}));
 }
 
+TEST(TNonblockingBatcherTest, IncompleteBatchAfterDeque)
+{
+    auto timeout = Quantum;
+    auto overTimeout = timeout * 2;
+
+    auto b = New<TNonblockingBatcher<int>>(TBatchSizeLimiter(2), timeout, false);
+    b->Enqueue(1);
+    b->Enqueue(2);
+    b->Enqueue(3);
+    auto e1 = b->DequeueBatch();
+    ASSERT_TRUE(e1.IsSet());
+    ASSERT_EQ(e1.Get().ValueOrThrow(), std::vector<int>({1, 2}));
+    Sleep(overTimeout);
+    b->Enqueue(4);
+    auto e2 = b->DequeueBatch();
+    ASSERT_TRUE(e2.IsSet());
+    ASSERT_EQ(e2.Get().ValueOrThrow(), std::vector<int>({3, 4}));
+}
+
 ////////////////////////////////////////////////////////////////////////////////
 
 } // namespace
diff --git a/yt/yt/library/profiling/impl.h b/yt/yt/library/profiling/impl.h
index 59f0b25ce7..48e863e882 100644
--- a/yt/yt/library/profiling/impl.h
+++ b/yt/yt/library/profiling/impl.h
@@ -55,7 +55,7 @@ public:
         const TTagSet& tags,
         TSensorOptions options) = 0;
 
-    virtual ITimerImplPtr RegisterTimeHistogram(
+    virtual IHistogramImplPtr RegisterTimeHistogram(
         const TString& name,
         const TTagSet& tags,
         TSensorOptions options) = 0;
diff --git a/yt/yt/library/profiling/public.h b/yt/yt/library/profiling/public.h
index db42520669..2abf46d968 100644
--- a/yt/yt/library/profiling/public.h
+++ b/yt/yt/library/profiling/public.h
@@ -25,6 +25,7 @@ DECLARE_REFCOUNTED_STRUCT(IHistogramImpl)
 DECLARE_REFCOUNTED_STRUCT(IRegistryImpl)
 DECLARE_REFCOUNTED_STRUCT(ISensorProducer)
 DECLARE_REFCOUNTED_CLASS(TBufferedProducer)
+DECLARE_REFCOUNTED_CLASS(THistogram)
 
 ////////////////////////////////////////////////////////////////////////////////
 
diff --git a/yt/yt/library/profiling/sensor.cpp b/yt/yt/library/profiling/sensor.cpp
index 560d9760bf..28bee90488 100644
--- a/yt/yt/library/profiling/sensor.cpp
+++ b/yt/yt/library/profiling/sensor.cpp
@@ -116,13 +116,16 @@ TSummary::operator bool() const
 
 ////////////////////////////////////////////////////////////////////////////////
 
-void TEventTimer::Record(TDuration value) const
+void TEventTimer::Record(TDuration value, int count) const
 {
-    if (!Timer_) {
-        return;
+    if (Histogram_) {
+        Histogram_->Add(value.SecondsFloat(), count);
+    } else if (Timer_) {
+        while (count > 0) {
+            Timer_->Record(value);
+            --count;
+        }
     }
-
-    Timer_->Record(value);
 }
 
 TEventTimer::operator bool() const
@@ -604,7 +607,7 @@ TEventTimer TProfiler::TimeHistogram(const TString& name, TDuration min, TDurati
     options.HistogramMax = max;
 
     TEventTimer timer;
-    timer.Timer_ = Impl_->RegisterTimeHistogram(Namespace_ + Prefix_ + name, Tags_, options);
+    timer.Histogram_ = Impl_->RegisterTimeHistogram(Namespace_ + Prefix_ + name, Tags_, options);
     return timer;
 }
 
@@ -617,7 +620,7 @@ TEventTimer TProfiler::TimeHistogram(const TString& name, std::vector<TDuration>
     TEventTimer timer;
     auto options = Options_;
     options.TimeHistogramBounds = std::move(bounds);
-    timer.Timer_ = Impl_->RegisterTimeHistogram(Namespace_ + Prefix_ + name, Tags_, options);
+    timer.Histogram_ = Impl_->RegisterTimeHistogram(Namespace_ + Prefix_ + name, Tags_, options);
     return timer;
 }
 
diff --git a/yt/yt/library/profiling/sensor.h b/yt/yt/library/profiling/sensor.h
index 2aa7299f32..f3b8c7dcf2 100644
--- a/yt/yt/library/profiling/sensor.h
+++ b/yt/yt/library/profiling/sensor.h
@@ -107,7 +107,7 @@ private:
 class TEventTimer
 {
 public:
-    void Record(TDuration value) const;
+    void Record(TDuration value, int count = 1) const;
 
     explicit operator bool() const;
 
@@ -115,6 +115,7 @@ private:
     friend class TProfiler;
 
     ITimerImplPtr Timer_;
+    IHistogramImplPtr Histogram_;
 };
 
 ////////////////////////////////////////////////////////////////////////////////
diff --git a/yt/yt/library/profiling/solomon/registry.cpp b/yt/yt/library/profiling/solomon/registry.cpp
index b231a64782..557e055e33 100644
--- a/yt/yt/library/profiling/solomon/registry.cpp
+++ b/yt/yt/library/profiling/solomon/registry.cpp
@@ -162,7 +162,7 @@ ITimerImplPtr TSolomonRegistry::RegisterTimerSummary(
         });
 }
 
-ITimerImplPtr TSolomonRegistry::RegisterTimeHistogram(
+IHistogramImplPtr TSolomonRegistry::RegisterTimeHistogram(
     const TString& name,
     const TTagSet& tags,
     TSensorOptions options)
diff --git a/yt/yt/library/profiling/solomon/registry.h b/yt/yt/library/profiling/solomon/registry.h
index 0569f18cd9..dbd684dfc3 100644
--- a/yt/yt/library/profiling/solomon/registry.h
+++ b/yt/yt/library/profiling/solomon/registry.h
@@ -80,7 +80,7 @@ public:
         const TTagSet& tags,
         TSensorOptions options) override;
 
-    ITimerImplPtr RegisterTimeHistogram(
+    IHistogramImplPtr RegisterTimeHistogram(
         const TString& name,
         const TTagSet& tags,
         TSensorOptions options) override;
diff --git a/yt/yt/library/profiling/solomon/sensor.h b/yt/yt/library/profiling/solomon/sensor.h
index 3ac491ef59..55d238821d 100644
--- a/yt/yt/library/profiling/solomon/sensor.h
+++ b/yt/yt/library/profiling/solomon/sensor.h
@@ -99,8 +99,6 @@ private:
 
 ////////////////////////////////////////////////////////////////////////////////
 
-DECLARE_REFCOUNTED_CLASS(THistogram)
-
 std::vector<double> GenerateGenericBucketBounds();
 
 class THistogram