From 084ff8d9934675912ee1af8687611eb34a630579 Mon Sep 17 00:00:00 2001
From: Alex Yarmula <ayarmula@fb.com>
Date: Thu, 19 Oct 2017 13:35:27 -0700
Subject: [PATCH] Handle timekeeperSingleton being nullptr in within()

Summary: When timekeeper singleton no longer exists during shutdown and folly::Singleton::try_get() can return nullptr, make sure nullptr is handled gracefully.

Reviewed By: yfeldblum

Differential Revision: D6101311

fbshipit-source-id: fefeddfbd048d1a7632688bb3526db15b685dd72
---
 folly/futures/Future-inl.h            |  8 ++++++--
 folly/futures/test/TimekeeperTest.cpp | 10 ++++++++++
 2 files changed, 16 insertions(+), 2 deletions(-)

diff --git a/folly/futures/Future-inl.h b/folly/futures/Future-inl.h
index 70a2c581..c7ddeb0a 100644
--- a/folly/futures/Future-inl.h
+++ b/folly/futures/Future-inl.h
@@ -1218,9 +1218,13 @@ Future<T> Future<T>::within(Duration dur, E e, Timekeeper* tk) {
   }
 
   std::shared_ptr<Timekeeper> tks;
-  if (!tk) {
+  if (LIKELY(!tk)) {
     tks = folly::detail::getTimekeeperSingleton();
-    tk = DCHECK_NOTNULL(tks.get());
+    tk = tks.get();
+  }
+
+  if (UNLIKELY(!tk)) {
+    return makeFuture<T>(NoTimekeeper());
   }
 
   auto ctx = std::make_shared<Context>(std::move(e));
diff --git a/folly/futures/test/TimekeeperTest.cpp b/folly/futures/test/TimekeeperTest.cpp
index 4567b93c..fdf17ed5 100644
--- a/folly/futures/test/TimekeeperTest.cpp
+++ b/folly/futures/test/TimekeeperTest.cpp
@@ -88,6 +88,16 @@ TEST(Timekeeper, futureSleepHandlesNullTimekeeperSingleton) {
   EXPECT_THROW(futures::sleep(one_ms).get(), NoTimekeeper);
 }
 
+TEST(Timekeeper, futureWithinHandlesNullTimekeeperSingleton) {
+  Singleton<ThreadWheelTimekeeper>::make_mock([] { return nullptr; });
+  SCOPE_EXIT {
+    Singleton<ThreadWheelTimekeeper>::make_mock();
+  };
+  Promise<int> p;
+  auto f = p.getFuture().within(one_ms);
+  EXPECT_THROW(f.get(), NoTimekeeper);
+}
+
 TEST(Timekeeper, futureDelayed) {
   auto t1 = now();
   auto dur = makeFuture()
-- 
2.34.1