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

1 files 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();
 }