Fix memory usage issues.

contrib/pzstd/utils/WorkQueue.h

@@ -99,6 +99,19 @@ class WorkQueue {
     return true;
   }
 
+  /**
+   * Sets the maximum queue size.  If `maxSize == 0` then it is unbounded.
+   *
+   * @param maxSize The new maximum queue size.
+   */
+  void setMaxSize(std::size_t maxSize) {
+    {
+      std::lock_guard<std::mutex> lock(mutex_);
+      maxSize_ = maxSize;
+    }
+    writerCv_.notify_all();
+  }
+
   /**
    * Promise that `push()` won't be called again, so once the queue is empty
    * there will never any more work.
@@ -149,6 +162,10 @@ class BufferWorkQueue {
     return result;
   }
 
+  void setMaxSize(std::size_t maxSize) {
+    queue_.setMaxSize(maxSize);
+  }
+
   void finish() {
     queue_.finish();
   }

contrib/pzstd/utils/test/WorkQueueTest.cpp

@@ -175,6 +175,27 @@ TEST(WorkQueue, BoundedSizePushAfterFinish) {
   pusher.join();
 }
 
+TEST(WorkQueue, SetMaxSize) {
+  WorkQueue<int> queue(2);
+  int result;
+  queue.push(5);
+  queue.push(6);
+  queue.setMaxSize(1);
+  std::thread pusher([&queue] {
+    queue.push(7);
+  });
+  // Dirtily try and make sure that pusher has run.
+  std::this_thread::sleep_for(std::chrono::seconds(1));
+  queue.finish();
+  EXPECT_TRUE(queue.pop(result));
+  EXPECT_EQ(5, result);
+  EXPECT_TRUE(queue.pop(result));
+  EXPECT_EQ(6, result);
+  EXPECT_FALSE(queue.pop(result));
+
+  pusher.join();
+}
+
 TEST(WorkQueue, BoundedSizeMPMC) {
   WorkQueue<int> queue(100);
   std::vector<int> results(10000, -1);