Do only one check for backpressure

Author: Rafał Hibner

cpp/src/arrow/dataset/dataset_writer.cc

@@ -54,7 +54,7 @@ class Throttle {
 
   bool Unthrottled() const { return max_value_ <= 0; }
 
-  Future<> Acquire(uint64_t values) {
+  std::optional<Future<>> Acquire(uint64_t values) {
     if (Unthrottled()) {
       return Future<>::MakeFinished();
     }
@@ -64,6 +64,8 @@ class Throttle {
       backpressure_ = Future<>::Make();
     } else {
       current_value_ += values;
+      DCHECK(backpressure_.is_finished());
+      return std::nullopt;
     }
     return backpressure_;
   }
@@ -72,9 +74,6 @@ class Throttle {
     if (Unthrottled()) {
       return;
     }
-    if(current_value_<values){
-      std::cout<<"Release: "<<current_value_<<"<"<<values<<std::endl;
-    }
     Future<> to_complete;
     {
       std::lock_guard<std::mutex> lg(mutex_);
@@ -666,7 +665,7 @@ class DatasetWriter::DatasetWriterImpl {
                                                        directory, prefix);
             }));
     std::shared_ptr<DatasetWriterDirectoryQueue> dir_queue = dir_queue_itr->second;
-    Future<> backpressure;
+    std::optional<Future<>> backpressure;
     while (batch) {
       // Keep opening new files until batch is done.
       std::shared_ptr<RecordBatch> remainder;
@@ -685,13 +684,13 @@ class DatasetWriter::DatasetWriterImpl {
       }
       backpressure =
           writer_state_->rows_in_flight_throttle.Acquire(next_chunk->num_rows());
-      if (!backpressure.is_finished()) {
+      if (backpressure) {
         EVENT_ON_CURRENT_SPAN("DatasetWriter::Backpressure::TooManyRowsQueued");
         break;
       }
       if (will_open_file) {
         backpressure = writer_state_->open_files_throttle.Acquire(1);
-        if (!backpressure.is_finished()) {
+        if (backpressure) {
           EVENT_ON_CURRENT_SPAN("DatasetWriter::Backpressure::TooManyOpenFiles");
           writer_state_->rows_in_flight_throttle.Release(next_chunk->num_rows());
           RETURN_NOT_OK(TryCloseLargestFile());
@@ -715,7 +714,8 @@ class DatasetWriter::DatasetWriterImpl {
     }
 
     if (batch) {
-      return backpressure.Then([this, batch, directory, prefix] {
+      DCHECK(backpressure);
+      return backpressure->Then([this, batch, directory, prefix] {
         return DoWriteRecordBatch(batch, directory, prefix);
       });
     }

cpp/src/arrow/dataset/dataset_writer_test.cc

@@ -280,32 +280,30 @@ TEST_F(DatasetWriterTestFixture, BatchGreaterThanMaxRowsQueued) {
 #pragma GCC push_options
 #pragma GCC optimize ("O0")
 TEST_F(DatasetWriterTestFixture, BatchWriteConcurrent) {
-  auto dataset_writer = MakeDatasetWriter(/*max_rows=*/200);
+  auto dataset_writer = MakeDatasetWriter(/*max_rows=*/5);
 
-  
-  for(int threads=20;threads>1;threads--){
-    for(int iter=2;iter<100;iter*=2){
-      for(int batch=2;batch<5000;batch*=2){
-        std::cout<<threads<<" "<<iter<<" "<<batch<<std::endl;
+
+  for(int threads=1;threads<5;threads++){
+    for(int iter=2;iter<=256;iter*=2){
+      for(int batch=2;batch<=64;batch*=2){
         std::vector<std::thread> workers;
         for(int i=0;i<threads;++i){
           workers.push_back(std::thread(
-          [&](){
+          [&,i=i](){
             for(int j=0;j<iter;++j){
               while(paused_){SleepABit();};
-              dataset_writer->WriteRecordBatch(MakeBatch(batch/threads), "");
+              dataset_writer->WriteRecordBatch(MakeBatch(batch+i+10*j), "");
             }
           }));
         }
         for (std::thread &t: workers) {
           if (t.joinable()) {
             t.join();
-          }      
+          }
           while(paused_){SleepABit();};
         }
       }
     }
-
   }
   EndWriterChecked(dataset_writer.get());
   ASSERT_EQ(paused_, false);