message: Fix flaky serial_device_batch_scheduler_test.
diff: @@ -145,8 +145,6 @@ TEST(SerialDeviceBatchSchedulerTest, PendingOnSerialDevice) {
  std::shared_ptr<SerialDeviceBatchScheduler<FakeTask>> scheduler;
  TF_ASSERT_OK(
  SerialDeviceBatchScheduler<FakeTask>::Create(options, &scheduler));
  - // Make sure batch processing thread has gone to sleep.
  - Env::Default()->SleepForMicroseconds(1000);
  int processed_batches = 0;
  Notification start_processing;
  auto queue_callback = [&mu, &processed_batches, &start_processing, &pending,
  @@ -161,28 +159,20 @@ TEST(SerialDeviceBatchSchedulerTest, PendingOnSerialDevice) {
  switch (batch_num) {
  case 1:
  start_processing.WaitForNotification();
  - {
  - mutex_lock l(mu);
  - pending = 2;
  - }
  - break;
  - case 2:
  - // No batches initially --> low traffic --> no adjustment.
  - CHECK_EQ(scheduler->in_flight_batches_limit(), 1);
  {
  mutex_lock l(mu);
  pending = 3;
  }
  break;
  - case 3:
  - // Pending at target --> no adjustment.
  + case 2:
  + // Either low traffic or pending at target --> no adjustment.
  CHECK_EQ(scheduler->in_flight_batches_limit(), 1);
  {
  mutex_lock l(mu);
  pending = 1;
  }
  break;
  - case 4:
  + case 3:
  // Small pending --> 2 additional threads added.
  CHECK_EQ(scheduler->in_flight_batches_limit(), 3);
  {
  @@ -196,8 +186,8 @@ TEST(SerialDeviceBatchSchedulerTest, PendingOnSerialDevice) {
  };
  std::unique_ptr<BatchScheduler<FakeTask>> queue;
  TF_ASSERT_OK(scheduler->AddQueue({}, queue_callback, &queue));
  - // Create 4 batches.
  - for (int i = 0; i < 4; i++) {
  + // Create 3 batches.
  + for (int i = 0; i < 3; i++) {
  TF_ASSERT_OK(ScheduleTask(800, queue.get()));
  }
  start_processing.Notify();
  @@ -295,12 +285,22 @@ TEST(SerialDeviceBatchSchedulerTest, DeleteQueue) {
  TF_ASSERT_OK(ScheduleTask(800, queue.get()));
  }
  std::unique_ptr<Thread> queue_deleter(Env::Default()->StartThread(
  - {}, "QueueDeleterThread", [&queue, &mu, &processed_batches] {
  + {}, "QueueDeleterThread",
  + [&queue, &mu, &processed_batches, scheduler]() mutable {
  // Delete queue, should be kept alive until empty.
  queue.reset();
  + {
  + mutex_lock l(mu);
  + // queue may be destroyed before 2nd batch finishes processing.
  + EXPECT_GT(processed_batches, 0);
  + }
  + // Delete scheduler, should be kept alive until all batches processed.
  + scheduler.reset();
  mutex_lock l(mu);
  EXPECT_EQ(processed_batches, 2);
  }));
  + // Release reference to scheduler, queue and callback above should keep alive.
  + scheduler.reset();
  // Give queue_deleter thread time to delete queue.
  Env::Default()->SleepForMicroseconds(1000);
  finish_processing.Notify();
  