TEST(NotificationQueueTest, ConsumeUntilDrainedStress) {
  for (size_t i = 0; i < 1 << 8; ++i) {
    // Basic tests: make sure we
    // - drain all the messages
    // - ignore any maxReadAtOnce
    // - can't add messages during draining
    EventBase eventBase;
    IntQueue queue;
    QueueConsumer consumer;
    consumer.fn = [&](int j) {
      EXPECT_THROW(queue.tryPutMessage(j), std::runtime_error);
      EXPECT_FALSE(queue.tryPutMessageNoThrow(j));
      EXPECT_THROW(queue.putMessage(j), std::runtime_error);
      std::vector<int> ints{1, 2, 3};
      EXPECT_THROW(
          queue.putMessages(ints.begin(), ints.end()),
          std::runtime_error);
    };
    consumer.setMaxReadAtOnce(10); // We should ignore this
    consumer.startConsuming(&eventBase, &queue);
    for (int j = 0; j < 20; j++) {
      queue.putMessage(j);
    }
    EXPECT_TRUE(consumer.consumeUntilDrained());
    EXPECT_EQ(20, consumer.messages.size());

    // Make sure there can only be one drainer at once
    folly::Baton<> callbackBaton, threadStartBaton;
    consumer.fn = [&](int /* i */) { callbackBaton.wait(); };
    QueueConsumer competingConsumer;
    competingConsumer.startConsuming(&eventBase, &queue);
    queue.putMessage(1);
    atomic<bool> raceA {false};
    atomic<bool> raceB {false};
    size_t numConsA = 0;
    size_t numConsB = 0;
    auto thread = std::thread([&]{
      threadStartBaton.post();
      raceB = consumer.consumeUntilDrained(&numConsB) && numConsB;
    });
    threadStartBaton.wait();
    raceA = competingConsumer.consumeUntilDrained(&numConsA) && numConsA;
    callbackBaton.post();
    thread.join();
    EXPECT_FALSE(raceA && raceB);
    EXPECT_TRUE(raceA || raceB);
    EXPECT_TRUE(raceA ^ raceB);
  }
}
void QueueTest::maxQueueSize() {
  // Create a queue with a maximum size of 5, and fill it up

  for (int n = 0; n < 5; ++n) {
    queue.tryPutMessage(n);
  }

  // Calling tryPutMessage() now should fail
  EXPECT_THROW(queue.tryPutMessage(5), std::overflow_error);

  EXPECT_FALSE(queue.tryPutMessageNoThrow(5));
  int val = 5;
  EXPECT_FALSE(queue.tryPutMessageNoThrow(std::move(val)));

  // Pop a message from the queue
  int result = -1;
  EXPECT_TRUE(queue.tryConsume(result));
  EXPECT_EQ(0, result);

  // We should be able to write another message now that we popped one off.
  queue.tryPutMessage(5);
  // But now we are full again.
  EXPECT_THROW(queue.tryPutMessage(6), std::overflow_error);
  // putMessage() should let us exceed the maximum
  queue.putMessage(6);

  // Pull another mesage off
  EXPECT_TRUE(queue.tryConsume(result));
  EXPECT_EQ(1, result);

  // tryPutMessage() should still fail since putMessage() actually put us over
  // the max.
  EXPECT_THROW(queue.tryPutMessage(7), std::overflow_error);

  // Pull another message off and try again
  EXPECT_TRUE(queue.tryConsume(result));
  EXPECT_EQ(2, result);
  queue.tryPutMessage(7);

  // Now pull all the remaining messages off
  EXPECT_TRUE(queue.tryConsume(result));
  EXPECT_EQ(3, result);
  EXPECT_TRUE(queue.tryConsume(result));
  EXPECT_EQ(4, result);
  EXPECT_TRUE(queue.tryConsume(result));
  EXPECT_EQ(5, result);
  EXPECT_TRUE(queue.tryConsume(result));
  EXPECT_EQ(6, result);
  EXPECT_TRUE(queue.tryConsume(result));
  EXPECT_EQ(7, result);

  // There should be no messages left
  result = -1;
  EXPECT_TRUE(!queue.tryConsume(result));
  EXPECT_EQ(-1, result);
}
void QueueTest::maxQueueSize() {
  // Create a queue with a maximum size of 5, and fill it up

  for (int n = 0; n < 5; ++n) {
    queue.tryPutMessage(n);
  }

  // Calling tryPutMessage() now should fail
  BOOST_CHECK_THROW(queue.tryPutMessage(5), TQueueFullException);

  BOOST_CHECK_EQUAL(queue.tryPutMessageNoThrow(5), false);
  int val = 5;
  BOOST_CHECK_EQUAL(queue.tryPutMessageNoThrow(std::move(val)), false);

  // Pop a message from the queue
  int result = -1;
  BOOST_CHECK(queue.tryConsume(result));
  BOOST_CHECK_EQUAL(result, 0);

  // We should be able to write another message now that we popped one off.
  queue.tryPutMessage(5);
  // But now we are full again.
  BOOST_CHECK_THROW(queue.tryPutMessage(6), TQueueFullException);
  // putMessage() should let us exceed the maximum
  queue.putMessage(6);

  // Pull another mesage off
  BOOST_CHECK(queue.tryConsume(result));
  BOOST_CHECK_EQUAL(result, 1);

  // tryPutMessage() should still fail since putMessage() actually put us over
  // the max.
  BOOST_CHECK_THROW(queue.tryPutMessage(7), TQueueFullException);

  // Pull another message off and try again
  BOOST_CHECK(queue.tryConsume(result));
  BOOST_CHECK_EQUAL(result, 2);
  queue.tryPutMessage(7);

  // Now pull all the remaining messages off
  BOOST_CHECK(queue.tryConsume(result));
  BOOST_CHECK_EQUAL(result, 3);
  BOOST_CHECK(queue.tryConsume(result));
  BOOST_CHECK_EQUAL(result, 4);
  BOOST_CHECK(queue.tryConsume(result));
  BOOST_CHECK_EQUAL(result, 5);
  BOOST_CHECK(queue.tryConsume(result));
  BOOST_CHECK_EQUAL(result, 6);
  BOOST_CHECK(queue.tryConsume(result));
  BOOST_CHECK_EQUAL(result, 7);

  // There should be no messages left
  result = -1;
  BOOST_CHECK(!queue.tryConsume(result));
  BOOST_CHECK_EQUAL(result, -1);
}
void QueueTest::fillQueue(bool expectFail) {
  try {
    for (uint32_t i = 0; i < 0x000fffff; i++) {
      queue.putMessage(i);
    }
    BOOST_CHECK(!expectFail);
  } catch (const apache::thrift::TLibraryException &ex) {
    BOOST_CHECK(expectFail);
  } catch (...) {
    BOOST_CHECK(false);
  }
}
void QueueTest::sendOne() {
  // Create a notification queue and a callback in this thread
  TEventBase eventBase;

  QueueConsumer consumer;
  consumer.fn = [&](int) {
    // Stop consuming after we receive 1 message
    consumer.stopConsuming();
  };
  consumer.startConsuming(&eventBase, &queue);

  // Start a new TEventBase thread to put a message on our queue
  ScopedEventBaseThread t1;
  t1.getEventBase()->runInEventBaseThread([&] {
    queue.putMessage(5);
  });

  // Loop until we receive the message
  eventBase.loop();

  const auto& messages = consumer.messages;
  BOOST_CHECK_EQUAL(messages.size(), 1);
  BOOST_CHECK_EQUAL(messages.at(0), 5);
}
/*
 * Test code that creates a NotificationQueue, then forks, and incorrectly
 * tries to send a message to the queue from the child process.
 *
 * The child process should crash in this scenario, since the child code has a
 * bug.  (Older versions of NotificationQueue didn't catch this in the child,
 * resulting in a crash in the parent process.)
 */
TEST(NotificationQueueTest, UseAfterFork) {
  IntQueue queue;
  int childStatus = 0;
  QueueConsumer consumer;

  // Boost sets a custom SIGCHLD handler, which fails the test if a child
  // process exits abnormally.  We don't want this.
  signal(SIGCHLD, SIG_DFL);

  // Log some info so users reading the test output aren't confused
  // by the child process' crash log messages.
  LOG(INFO) << "This test makes sure the child process crashes.  "
    << "Error log messagges and a backtrace are expected.";

  {
    // Start a separate thread consuming from the queue
    ScopedEventBaseThread t1;
    t1.getEventBase()->runInEventBaseThread([&] {
      consumer.startConsuming(t1.getEventBase(), &queue);
    });

    // Send a message to it, just for sanity checking
    queue.putMessage(1234);

    // Fork
    pid_t pid = fork();
    if (pid == 0) {
      // The boost test framework installs signal handlers to catch errors.
      // We only want to catch in the parent.  In the child let SIGABRT crash
      // us normally.
      signal(SIGABRT, SIG_DFL);

      // Child.
      // We're horrible people, so we try to send a message to the queue
      // that is being consumed in the parent process.
      //
      // The putMessage() call should catch this error, and crash our process.
      queue.putMessage(9876);
      // We shouldn't reach here.
      _exit(0);
    }
    PCHECK(pid > 0);

    // Parent.  Wait for the child to exit.
    auto waited = waitpid(pid, &childStatus, 0);
    EXPECT_EQ(pid, waited);

    // Send another message to the queue before we terminate the thread.
    queue.putMessage(5678);
  }

  // The child process should have crashed when it tried to call putMessage()
  // on our NotificationQueue.
  EXPECT_TRUE(WIFSIGNALED(childStatus));
  EXPECT_EQ(SIGABRT, WTERMSIG(childStatus));

  // Make sure the parent saw the expected messages.
  // It should have gotten 1234 and 5678 from the parent process, but not
  // 9876 from the child.
  EXPECT_EQ(2, consumer.messages.size());
  EXPECT_EQ(1234, consumer.messages.front());
  consumer.messages.pop_front();
  EXPECT_EQ(5678, consumer.messages.front());
  consumer.messages.pop_front();
}
void QueueTest::maxReadAtOnce() {
  // Add 100 messages to the queue
  for (int n = 0; n < 100; ++n) {
    queue.putMessage(n);
  }

  EventBase eventBase;

  // Record how many messages were processed each loop iteration.
  uint32_t messagesThisLoop = 0;
  std::vector<uint32_t> messagesPerLoop;
  std::function<void()> loopFinished = [&] {
    // Record the current number of messages read this loop
    messagesPerLoop.push_back(messagesThisLoop);
    // Reset messagesThisLoop to 0 for the next loop
    messagesThisLoop = 0;

    // To prevent use-after-free bugs when eventBase destructs,
    // prevent calling runInLoop any more after the test is finished.
    // 55 == number of times loop should run.
    if (messagesPerLoop.size() != 55) {
      // Reschedule ourself to run at the end of the next loop
      eventBase.runInLoop(loopFinished);
    }
  };
  // Schedule the first call to loopFinished
  eventBase.runInLoop(loopFinished);

  QueueConsumer consumer;
  // Read the first 50 messages 10 at a time.
  consumer.setMaxReadAtOnce(10);
  consumer.fn = [&](int value) {
    ++messagesThisLoop;
    // After 50 messages, drop to reading only 1 message at a time.
    if (value == 50) {
      consumer.setMaxReadAtOnce(1);
    }
    // Terminate the loop when we reach the end of the messages.
    if (value == 99) {
      eventBase.terminateLoopSoon();
    }
  };
  consumer.startConsuming(&eventBase, &queue);

  // Run the event loop until the consumer terminates it
  eventBase.loop();

  // The consumer should have read all 100 messages in order
  EXPECT_EQ(100, consumer.messages.size());
  for (int n = 0; n < 100; ++n) {
    EXPECT_EQ(n, consumer.messages.at(n));
  }

  // Currently EventBase happens to still run the loop callbacks even after
  // terminateLoopSoon() is called.  However, we don't really want to depend on
  // this behavior.  In case this ever changes in the future, add
  // messagesThisLoop to messagesPerLoop in loop callback isn't invoked for the
  // last loop iteration.
  if (messagesThisLoop > 0) {
    messagesPerLoop.push_back(messagesThisLoop);
    messagesThisLoop = 0;
  }

  // For the first 5 loops it should have read 10 messages each time.
  // After that it should have read 1 messages per loop for the next 50 loops.
  EXPECT_EQ(55, messagesPerLoop.size());
  for (int n = 0; n < 5; ++n) {
    EXPECT_EQ(10, messagesPerLoop.at(n));
  }
  for (int n = 5; n < 55; ++n) {
    EXPECT_EQ(1, messagesPerLoop.at(n));
  }
}
void QueueTest::multiConsumer() {
  uint32_t numConsumers = 8;
  uint32_t numMessages = 10000;

  // Create several consumers each running in their own EventBase thread
  vector<QueueConsumer> consumers(numConsumers);
  vector<ScopedEventBaseThread> threads(numConsumers);

  for (uint32_t consumerIdx = 0; consumerIdx < numConsumers; ++consumerIdx) {
    QueueConsumer* consumer = &consumers[consumerIdx];

    consumer->fn = [consumer, consumerIdx, this](int value) {
      // Treat 0 as a signal to stop.
      if (value == 0) {
        consumer->stopConsuming();
        // Put a message on the terminationQueue to indicate we have stopped
        terminationQueue.putMessage(consumerIdx);
      }
    };

    EventBase* eventBase = threads[consumerIdx].getEventBase();
    eventBase->runInEventBaseThread([eventBase, consumer, this] {
      consumer->startConsuming(eventBase, &queue);
    });
  }

  // Now add a number of messages from this thread
  // Start at 1 rather than 0, since 0 is the signal to stop.
  for (uint32_t n = 1; n < numMessages; ++n) {
    queue.putMessage(n);
  }
  // Now add a 0 for each consumer, to signal them to stop
  for (uint32_t n = 0; n < numConsumers; ++n) {
    queue.putMessage(0);
  }

  // Wait until we get notified that all of the consumers have stopped
  // We use a separate notification queue for this.
  QueueConsumer terminationConsumer;
  vector<uint32_t> consumersStopped(numConsumers, 0);
  uint32_t consumersRemaining = numConsumers;
  terminationConsumer.fn = [&](int consumerIdx) {
    --consumersRemaining;
    if (consumersRemaining == 0) {
      terminationConsumer.stopConsuming();
    }

    EXPECT_GE(consumerIdx, 0);
    EXPECT_LT(consumerIdx, numConsumers);
    ++consumersStopped[consumerIdx];
  };
  EventBase eventBase;
  terminationConsumer.startConsuming(&eventBase, &terminationQueue);
  eventBase.loop();

  // Verify that we saw exactly 1 stop message for each consumer
  for (uint32_t n = 0; n < numConsumers; ++n) {
    EXPECT_EQ(1, consumersStopped[n]);
  }

  // Validate that every message sent to the main queue was received exactly
  // once.
  vector<int> messageCount(numMessages, 0);
  for (uint32_t n = 0; n < numConsumers; ++n) {
    for (int msg : consumers[n].messages) {
      EXPECT_GE(msg, 0);
      EXPECT_LT(msg, numMessages);
      ++messageCount[msg];
    }
  }

  // 0 is the signal to stop, and should have been received once by each
  // consumer
  EXPECT_EQ(numConsumers, messageCount[0]);
  // All other messages should have been received exactly once
  for (uint32_t n = 1; n < numMessages; ++n) {
    EXPECT_EQ(1, messageCount[n]);
  }
}