From 48631f367ebcbb23f6fe73482398b594279c2dea Mon Sep 17 00:00:00 2001
From: "A. Unique TensorFlower" <gardener@tensorflow.org>
Date: Mon, 6 Nov 2017 15:09:23 -0800
Subject: Fix race bug in AdaptiveSharedBatchScheduler.

In ASBSQueue::Schedule, when a new batch is created, it was added to the scheduler outside of the queue's lock.  This was done to prevent any unforeseen interactions between the queue lock and scheduler lock.  However, this wasn't being done in a thread safe way.

PiperOrigin-RevId: 174769383
---
 tensorflow/contrib/batching/adaptive_shared_batch_scheduler.h | 7 +++----
 1 file changed, 3 insertions(+), 4 deletions(-)

diff --git a/tensorflow/contrib/batching/adaptive_shared_batch_scheduler.h b/tensorflow/contrib/batching/adaptive_shared_batch_scheduler.h
index a0606427a5..6ed177e001 100644
--- a/tensorflow/contrib/batching/adaptive_shared_batch_scheduler.h
+++ b/tensorflow/contrib/batching/adaptive_shared_batch_scheduler.h
@@ -399,7 +399,7 @@ ASBSQueue<TaskType>::~ASBSQueue() {
 
 template <typename TaskType>
 Status ASBSQueue<TaskType>::Schedule(std::unique_ptr<TaskType>* task) {
-  bool added_new_batch = false;
+  ASBSBatch<TaskType>* new_batch = nullptr;
   size_t size = (*task)->size();
   if (size > options_.max_batch_size) {
     return errors::InvalidArgument("Task size ", size,
@@ -418,15 +418,14 @@ Status ASBSQueue<TaskType>::Schedule(std::unique_ptr<TaskType>* task) {
       current_batch_ = nullptr;
     }
     if (!current_batch_) {
-      added_new_batch = true;
       num_enqueued_batches_++;
-      current_batch_ =
+      current_batch_ = new_batch =
           new ASBSBatch<TaskType>(this, scheduler_->GetEnv()->NowMicros());
     }
     current_batch_->AddTask(std::move(*task));
     num_enqueued_tasks_++;
   }
-  if (added_new_batch) scheduler_->AddBatch(current_batch_);
+  if (new_batch != nullptr) scheduler_->AddBatch(new_batch);
   return Status::OK();
 }
 
-- 
cgit v1.2.3