folly/fibers/test/FibersTest.cpp

   1 /*
   2  * Copyright 2017-present Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16 #include <atomic>
  17 #include <thread>
  18 #include <vector>
  19
  20 #include <folly/Memory.h>
  21 #include <folly/Random.h>
  22 #include <folly/futures/Future.h>
  23
  24 #include <folly/Conv.h>
  25 #include <folly/fibers/AddTasks.h>
  26 #include <folly/fibers/AtomicBatchDispatcher.h>
  27 #include <folly/fibers/BatchDispatcher.h>
  28 #include <folly/fibers/EventBaseLoopController.h>
  29 #include <folly/fibers/FiberManager.h>
  30 #include <folly/fibers/FiberManagerMap.h>
  31 #include <folly/fibers/GenericBaton.h>
  32 #include <folly/fibers/Semaphore.h>
  33 #include <folly/fibers/SimpleLoopController.h>
  34 #include <folly/fibers/TimedMutex.h>
  35 #include <folly/fibers/WhenN.h>
  36 #include <folly/io/async/ScopedEventBaseThread.h>
  37 #include <folly/portability/GTest.h>
  38
  39 using namespace folly::fibers;
  40
  41 using folly::Try;
  42
  43 TEST(FiberManager, batonTimedWaitTimeout) {
  44   bool taskAdded = false;
  45   size_t iterations = 0;
  46
  47   FiberManager manager(std::make_unique<SimpleLoopController>());
  48   auto& loopController =
  49       dynamic_cast<SimpleLoopController&>(manager.loopController());
  50
  51   auto loopFunc = [&]() {
  52     if (!taskAdded) {
  53       manager.addTask([&]() {
  54         Baton baton;
  55
  56         auto res = baton.try_wait_for(std::chrono::milliseconds(230));
  57
  58         EXPECT_FALSE(res);
  59         EXPECT_EQ(5, iterations);
  60
  61         loopController.stop();
  62       });
  63       manager.addTask([&]() {
  64         Baton baton;
  65
  66         auto res = baton.try_wait_for(std::chrono::milliseconds(130));
  67
  68         EXPECT_FALSE(res);
  69         EXPECT_EQ(3, iterations);
  70
  71         loopController.stop();
  72       });
  73       taskAdded = true;
  74     } else {
  75       std::this_thread::sleep_for(std::chrono::milliseconds(50));
  76       iterations++;
  77     }
  78   };
  79
  80   loopController.loop(std::move(loopFunc));
  81 }
  82
  83 TEST(FiberManager, batonTimedWaitPost) {
  84   bool taskAdded = false;
  85   size_t iterations = 0;
  86   Baton* baton_ptr;
  87
  88   FiberManager manager(std::make_unique<SimpleLoopController>());
  89   auto& loopController =
  90       dynamic_cast<SimpleLoopController&>(manager.loopController());
  91
  92   auto loopFunc = [&]() {
  93     if (!taskAdded) {
  94       manager.addTask([&]() {
  95         Baton baton;
  96         baton_ptr = &baton;
  97
  98         auto res = baton.try_wait_for(std::chrono::milliseconds(130));
  99
 100         EXPECT_TRUE(res);
 101         EXPECT_EQ(2, iterations);
 102
 103         loopController.stop();
 104       });
 105       taskAdded = true;
 106     } else {
 107       std::this_thread::sleep_for(std::chrono::milliseconds(50));
 108       iterations++;
 109       if (iterations == 2) {
 110         baton_ptr->post();
 111       }
 112     }
 113   };
 114
 115   loopController.loop(std::move(loopFunc));
 116 }
 117
 118 TEST(FiberManager, batonTimedWaitTimeoutEvb) {
 119   size_t tasksComplete = 0;
 120
 121   folly::EventBase evb;
 122
 123   FiberManager manager(std::make_unique<EventBaseLoopController>());
 124   dynamic_cast<EventBaseLoopController&>(manager.loopController())
 125       .attachEventBase(evb);
 126
 127   auto task = [&](size_t timeout_ms) {
 128     Baton baton;
 129
 130     auto start = EventBaseLoopController::Clock::now();
 131     auto res = baton.try_wait_for(std::chrono::milliseconds(timeout_ms));
 132     auto finish = EventBaseLoopController::Clock::now();
 133
 134     EXPECT_FALSE(res);
 135
 136     auto duration_ms =
 137         std::chrono::duration_cast<std::chrono::milliseconds>(finish - start);
 138
 139     EXPECT_GT(duration_ms.count(), timeout_ms - 50);
 140     EXPECT_LT(duration_ms.count(), timeout_ms + 50);
 141
 142     if (++tasksComplete == 2) {
 143       evb.terminateLoopSoon();
 144     }
 145   };
 146
 147   evb.runInEventBaseThread([&]() {
 148     manager.addTask([&]() { task(500); });
 149     manager.addTask([&]() { task(250); });
 150   });
 151
 152   evb.loopForever();
 153
 154   EXPECT_EQ(2, tasksComplete);
 155 }
 156
 157 TEST(FiberManager, batonTimedWaitPostEvb) {
 158   size_t tasksComplete = 0;
 159
 160   folly::EventBase evb;
 161
 162   FiberManager manager(std::make_unique<EventBaseLoopController>());
 163   dynamic_cast<EventBaseLoopController&>(manager.loopController())
 164       .attachEventBase(evb);
 165
 166   evb.runInEventBaseThread([&]() {
 167     manager.addTask([&]() {
 168       Baton baton;
 169
 170       evb.tryRunAfterDelay([&]() { baton.post(); }, 100);
 171
 172       auto start = EventBaseLoopController::Clock::now();
 173       auto res = baton.try_wait_for(std::chrono::milliseconds(130));
 174       auto finish = EventBaseLoopController::Clock::now();
 175
 176       EXPECT_TRUE(res);
 177
 178       auto duration_ms =
 179           std::chrono::duration_cast<std::chrono::milliseconds>(finish - start);
 180
 181       EXPECT_TRUE(duration_ms.count() > 95 && duration_ms.count() < 110);
 182
 183       if (++tasksComplete == 1) {
 184         evb.terminateLoopSoon();
 185       }
 186     });
 187   });
 188
 189   evb.loopForever();
 190
 191   EXPECT_EQ(1, tasksComplete);
 192 }
 193
 194 TEST(FiberManager, batonTryWait) {
 195   FiberManager manager(std::make_unique<SimpleLoopController>());
 196
 197   // Check if try_wait and post work as expected
 198   Baton b;
 199
 200   manager.addTask([&]() {
 201     while (!b.try_wait()) {
 202     }
 203   });
 204   auto thr = std::thread([&]() {
 205     std::this_thread::sleep_for(std::chrono::milliseconds(300));
 206     b.post();
 207   });
 208
 209   manager.loopUntilNoReady();
 210   thr.join();
 211
 212   Baton c;
 213
 214   // Check try_wait without post
 215   manager.addTask([&]() {
 216     int cnt = 100;
 217     while (cnt && !c.try_wait()) {
 218       cnt--;
 219     }
 220     EXPECT_TRUE(!c.try_wait()); // must still hold
 221     EXPECT_EQ(cnt, 0);
 222   });
 223
 224   manager.loopUntilNoReady();
 225 }
 226
 227 TEST(FiberManager, genericBatonFiberWait) {
 228   FiberManager manager(std::make_unique<SimpleLoopController>());
 229
 230   GenericBaton b;
 231   bool fiberRunning = false;
 232
 233   manager.addTask([&]() {
 234     EXPECT_EQ(manager.hasActiveFiber(), true);
 235     fiberRunning = true;
 236     b.wait();
 237     fiberRunning = false;
 238   });
 239
 240   EXPECT_FALSE(fiberRunning);
 241   manager.loopUntilNoReady();
 242   EXPECT_TRUE(fiberRunning); // ensure fiber still active
 243
 244   auto thr = std::thread([&]() {
 245     std::this_thread::sleep_for(std::chrono::milliseconds(300));
 246     b.post();
 247   });
 248
 249   while (fiberRunning) {
 250     manager.loopUntilNoReady();
 251   }
 252
 253   thr.join();
 254 }
 255
 256 TEST(FiberManager, genericBatonThreadWait) {
 257   FiberManager manager(std::make_unique<SimpleLoopController>());
 258   GenericBaton b;
 259   std::atomic<bool> threadWaiting(false);
 260
 261   auto thr = std::thread([&]() {
 262     threadWaiting = true;
 263     b.wait();
 264     threadWaiting = false;
 265   });
 266
 267   while (!threadWaiting) {
 268   }
 269   std::this_thread::sleep_for(std::chrono::milliseconds(300));
 270
 271   manager.addTask([&]() {
 272     EXPECT_EQ(manager.hasActiveFiber(), true);
 273     EXPECT_TRUE(threadWaiting);
 274     b.post();
 275     while (threadWaiting) {
 276     }
 277   });
 278
 279   manager.loopUntilNoReady();
 280   thr.join();
 281 }
 282
 283 TEST(FiberManager, addTasksNoncopyable) {
 284   std::vector<Promise<int>> pendingFibers;
 285   bool taskAdded = false;
 286
 287   FiberManager manager(std::make_unique<SimpleLoopController>());
 288   auto& loopController =
 289       dynamic_cast<SimpleLoopController&>(manager.loopController());
 290
 291   auto loopFunc = [&]() {
 292     if (!taskAdded) {
 293       manager.addTask([&]() {
 294         std::vector<std::function<std::unique_ptr<int>()>> funcs;
 295         for (int i = 0; i < 3; ++i) {
 296           funcs.push_back([i, &pendingFibers]() {
 297             await([&pendingFibers](Promise<int> promise) {
 298               pendingFibers.push_back(std::move(promise));
 299             });
 300             return std::make_unique<int>(i * 2 + 1);
 301           });
 302         }
 303
 304         auto iter = addTasks(funcs.begin(), funcs.end());
 305
 306         size_t n = 0;
 307         while (iter.hasNext()) {
 308           auto result = iter.awaitNext();
 309           EXPECT_EQ(2 * iter.getTaskID() + 1, *result);
 310           EXPECT_GE(2 - n, pendingFibers.size());
 311           ++n;
 312         }
 313         EXPECT_EQ(3, n);
 314       });
 315       taskAdded = true;
 316     } else if (pendingFibers.size()) {
 317       pendingFibers.back().setValue(0);
 318       pendingFibers.pop_back();
 319     } else {
 320       loopController.stop();
 321     }
 322   };
 323
 324   loopController.loop(std::move(loopFunc));
 325 }
 326
 327 TEST(FiberManager, awaitThrow) {
 328   folly::EventBase evb;
 329   struct ExpectedException {};
 330   getFiberManager(evb)
 331       .addTaskFuture([&] {
 332         EXPECT_THROW(
 333             await([](Promise<int> p) {
 334               p.setValue(42);
 335               throw ExpectedException();
 336             }),
 337             ExpectedException);
 338
 339         EXPECT_THROW(
 340             await([&](Promise<int> p) {
 341               evb.runInEventBaseThread([p = std::move(p)]() mutable {
 342                 p.setValue(42);
 343               });
 344               throw ExpectedException();
 345             }),
 346             ExpectedException);
 347       })
 348       .waitVia(&evb);
 349 }
 350
 351 TEST(FiberManager, addTasksThrow) {
 352   std::vector<Promise<int>> pendingFibers;
 353   bool taskAdded = false;
 354
 355   FiberManager manager(std::make_unique<SimpleLoopController>());
 356   auto& loopController =
 357       dynamic_cast<SimpleLoopController&>(manager.loopController());
 358
 359   auto loopFunc = [&]() {
 360     if (!taskAdded) {
 361       manager.addTask([&]() {
 362         std::vector<std::function<int()>> funcs;
 363         for (size_t i = 0; i < 3; ++i) {
 364           funcs.push_back([i, &pendingFibers]() {
 365             await([&pendingFibers](Promise<int> promise) {
 366               pendingFibers.push_back(std::move(promise));
 367             });
 368             if (i % 2 == 0) {
 369               throw std::runtime_error("Runtime");
 370             }
 371             return i * 2 + 1;
 372           });
 373         }
 374
 375         auto iter = addTasks(funcs.begin(), funcs.end());
 376
 377         size_t n = 0;
 378         while (iter.hasNext()) {
 379           try {
 380             int result = iter.awaitNext();
 381             EXPECT_EQ(1, iter.getTaskID() % 2);
 382             EXPECT_EQ(2 * iter.getTaskID() + 1, result);
 383           } catch (...) {
 384             EXPECT_EQ(0, iter.getTaskID() % 2);
 385           }
 386           EXPECT_GE(2 - n, pendingFibers.size());
 387           ++n;
 388         }
 389         EXPECT_EQ(3, n);
 390       });
 391       taskAdded = true;
 392     } else if (pendingFibers.size()) {
 393       pendingFibers.back().setValue(0);
 394       pendingFibers.pop_back();
 395     } else {
 396       loopController.stop();
 397     }
 398   };
 399
 400   loopController.loop(std::move(loopFunc));
 401 }
 402
 403 TEST(FiberManager, addTasksVoid) {
 404   std::vector<Promise<int>> pendingFibers;
 405   bool taskAdded = false;
 406
 407   FiberManager manager(std::make_unique<SimpleLoopController>());
 408   auto& loopController =
 409       dynamic_cast<SimpleLoopController&>(manager.loopController());
 410
 411   auto loopFunc = [&]() {
 412     if (!taskAdded) {
 413       manager.addTask([&]() {
 414         std::vector<std::function<void()>> funcs;
 415         for (size_t i = 0; i < 3; ++i) {
 416           funcs.push_back([&pendingFibers]() {
 417             await([&pendingFibers](Promise<int> promise) {
 418               pendingFibers.push_back(std::move(promise));
 419             });
 420           });
 421         }
 422
 423         auto iter = addTasks(funcs.begin(), funcs.end());
 424
 425         size_t n = 0;
 426         while (iter.hasNext()) {
 427           iter.awaitNext();
 428           EXPECT_GE(2 - n, pendingFibers.size());
 429           ++n;
 430         }
 431         EXPECT_EQ(3, n);
 432       });
 433       taskAdded = true;
 434     } else if (pendingFibers.size()) {
 435       pendingFibers.back().setValue(0);
 436       pendingFibers.pop_back();
 437     } else {
 438       loopController.stop();
 439     }
 440   };
 441
 442   loopController.loop(std::move(loopFunc));
 443 }
 444
 445 TEST(FiberManager, addTasksVoidThrow) {
 446   std::vector<Promise<int>> pendingFibers;
 447   bool taskAdded = false;
 448
 449   FiberManager manager(std::make_unique<SimpleLoopController>());
 450   auto& loopController =
 451       dynamic_cast<SimpleLoopController&>(manager.loopController());
 452
 453   auto loopFunc = [&]() {
 454     if (!taskAdded) {
 455       manager.addTask([&]() {
 456         std::vector<std::function<void()>> funcs;
 457         for (size_t i = 0; i < 3; ++i) {
 458           funcs.push_back([i, &pendingFibers]() {
 459             await([&pendingFibers](Promise<int> promise) {
 460               pendingFibers.push_back(std::move(promise));
 461             });
 462             if (i % 2 == 0) {
 463               throw std::runtime_error("");
 464             }
 465           });
 466         }
 467
 468         auto iter = addTasks(funcs.begin(), funcs.end());
 469
 470         size_t n = 0;
 471         while (iter.hasNext()) {
 472           try {
 473             iter.awaitNext();
 474             EXPECT_EQ(1, iter.getTaskID() % 2);
 475           } catch (...) {
 476             EXPECT_EQ(0, iter.getTaskID() % 2);
 477           }
 478           EXPECT_GE(2 - n, pendingFibers.size());
 479           ++n;
 480         }
 481         EXPECT_EQ(3, n);
 482       });
 483       taskAdded = true;
 484     } else if (pendingFibers.size()) {
 485       pendingFibers.back().setValue(0);
 486       pendingFibers.pop_back();
 487     } else {
 488       loopController.stop();
 489     }
 490   };
 491
 492   loopController.loop(std::move(loopFunc));
 493 }
 494
 495 TEST(FiberManager, addTasksReserve) {
 496   std::vector<Promise<int>> pendingFibers;
 497   bool taskAdded = false;
 498
 499   FiberManager manager(std::make_unique<SimpleLoopController>());
 500   auto& loopController =
 501       dynamic_cast<SimpleLoopController&>(manager.loopController());
 502
 503   auto loopFunc = [&]() {
 504     if (!taskAdded) {
 505       manager.addTask([&]() {
 506         std::vector<std::function<void()>> funcs;
 507         for (size_t i = 0; i < 3; ++i) {
 508           funcs.push_back([&pendingFibers]() {
 509             await([&pendingFibers](Promise<int> promise) {
 510               pendingFibers.push_back(std::move(promise));
 511             });
 512           });
 513         }
 514
 515         auto iter = addTasks(funcs.begin(), funcs.end());
 516
 517         iter.reserve(2);
 518         EXPECT_TRUE(iter.hasCompleted());
 519         EXPECT_TRUE(iter.hasPending());
 520         EXPECT_TRUE(iter.hasNext());
 521
 522         iter.awaitNext();
 523         EXPECT_TRUE(iter.hasCompleted());
 524         EXPECT_TRUE(iter.hasPending());
 525         EXPECT_TRUE(iter.hasNext());
 526
 527         iter.awaitNext();
 528         EXPECT_FALSE(iter.hasCompleted());
 529         EXPECT_TRUE(iter.hasPending());
 530         EXPECT_TRUE(iter.hasNext());
 531
 532         iter.awaitNext();
 533         EXPECT_FALSE(iter.hasCompleted());
 534         EXPECT_FALSE(iter.hasPending());
 535         EXPECT_FALSE(iter.hasNext());
 536       });
 537       taskAdded = true;
 538     } else if (pendingFibers.size()) {
 539       pendingFibers.back().setValue(0);
 540       pendingFibers.pop_back();
 541     } else {
 542       loopController.stop();
 543     }
 544   };
 545
 546   loopController.loop(std::move(loopFunc));
 547 }
 548
 549 TEST(FiberManager, addTaskDynamic) {
 550   folly::EventBase evb;
 551
 552   Baton batons[3];
 553
 554   auto makeTask = [&](size_t taskId) {
 555     return [&, taskId]() -> size_t {
 556       batons[taskId].wait();
 557       return taskId;
 558     };
 559   };
 560
 561   getFiberManager(evb)
 562       .addTaskFuture([&]() {
 563         TaskIterator<size_t> iterator;
 564
 565         iterator.addTask(makeTask(0));
 566         iterator.addTask(makeTask(1));
 567
 568         batons[1].post();
 569
 570         EXPECT_EQ(1, iterator.awaitNext());
 571
 572         iterator.addTask(makeTask(2));
 573
 574         batons[2].post();
 575
 576         EXPECT_EQ(2, iterator.awaitNext());
 577
 578         batons[0].post();
 579
 580         EXPECT_EQ(0, iterator.awaitNext());
 581       })
 582       .waitVia(&evb);
 583 }
 584
 585 TEST(FiberManager, forEach) {
 586   std::vector<Promise<int>> pendingFibers;
 587   bool taskAdded = false;
 588
 589   FiberManager manager(std::make_unique<SimpleLoopController>());
 590   auto& loopController =
 591       dynamic_cast<SimpleLoopController&>(manager.loopController());
 592
 593   auto loopFunc = [&]() {
 594     if (!taskAdded) {
 595       manager.addTask([&]() {
 596         std::vector<std::function<int()>> funcs;
 597         for (size_t i = 0; i < 3; ++i) {
 598           funcs.push_back([i, &pendingFibers]() {
 599             await([&pendingFibers](Promise<int> promise) {
 600               pendingFibers.push_back(std::move(promise));
 601             });
 602             return i * 2 + 1;
 603           });
 604         }
 605
 606         std::vector<std::pair<size_t, int>> results;
 607         forEach(funcs.begin(), funcs.end(), [&results](size_t id, int result) {
 608           results.emplace_back(id, result);
 609         });
 610         EXPECT_EQ(3, results.size());
 611         EXPECT_TRUE(pendingFibers.empty());
 612         for (size_t i = 0; i < 3; ++i) {
 613           EXPECT_EQ(results[i].first * 2 + 1, results[i].second);
 614         }
 615       });
 616       taskAdded = true;
 617     } else if (pendingFibers.size()) {
 618       pendingFibers.back().setValue(0);
 619       pendingFibers.pop_back();
 620     } else {
 621       loopController.stop();
 622     }
 623   };
 624
 625   loopController.loop(std::move(loopFunc));
 626 }
 627
 628 TEST(FiberManager, collectN) {
 629   std::vector<Promise<int>> pendingFibers;
 630   bool taskAdded = false;
 631
 632   FiberManager manager(std::make_unique<SimpleLoopController>());
 633   auto& loopController =
 634       dynamic_cast<SimpleLoopController&>(manager.loopController());
 635
 636   auto loopFunc = [&]() {
 637     if (!taskAdded) {
 638       manager.addTask([&]() {
 639         std::vector<std::function<int()>> funcs;
 640         for (size_t i = 0; i < 3; ++i) {
 641           funcs.push_back([i, &pendingFibers]() {
 642             await([&pendingFibers](Promise<int> promise) {
 643               pendingFibers.push_back(std::move(promise));
 644             });
 645             return i * 2 + 1;
 646           });
 647         }
 648
 649         auto results = collectN(funcs.begin(), funcs.end(), 2);
 650         EXPECT_EQ(2, results.size());
 651         EXPECT_EQ(1, pendingFibers.size());
 652         for (size_t i = 0; i < 2; ++i) {
 653           EXPECT_EQ(results[i].first * 2 + 1, results[i].second);
 654         }
 655       });
 656       taskAdded = true;
 657     } else if (pendingFibers.size()) {
 658       pendingFibers.back().setValue(0);
 659       pendingFibers.pop_back();
 660     } else {
 661       loopController.stop();
 662     }
 663   };
 664
 665   loopController.loop(std::move(loopFunc));
 666 }
 667
 668 TEST(FiberManager, collectNThrow) {
 669   std::vector<Promise<int>> pendingFibers;
 670   bool taskAdded = false;
 671
 672   FiberManager manager(std::make_unique<SimpleLoopController>());
 673   auto& loopController =
 674       dynamic_cast<SimpleLoopController&>(manager.loopController());
 675
 676   auto loopFunc = [&]() {
 677     if (!taskAdded) {
 678       manager.addTask([&]() {
 679         std::vector<std::function<int()>> funcs;
 680         for (size_t i = 0; i < 3; ++i) {
 681           funcs.push_back([&pendingFibers]() -> size_t {
 682             await([&pendingFibers](Promise<int> promise) {
 683               pendingFibers.push_back(std::move(promise));
 684             });
 685             throw std::runtime_error("Runtime");
 686           });
 687         }
 688
 689         try {
 690           collectN(funcs.begin(), funcs.end(), 2);
 691         } catch (...) {
 692           EXPECT_EQ(1, pendingFibers.size());
 693         }
 694       });
 695       taskAdded = true;
 696     } else if (pendingFibers.size()) {
 697       pendingFibers.back().setValue(0);
 698       pendingFibers.pop_back();
 699     } else {
 700       loopController.stop();
 701     }
 702   };
 703
 704   loopController.loop(std::move(loopFunc));
 705 }
 706
 707 TEST(FiberManager, collectNVoid) {
 708   std::vector<Promise<int>> pendingFibers;
 709   bool taskAdded = false;
 710
 711   FiberManager manager(std::make_unique<SimpleLoopController>());
 712   auto& loopController =
 713       dynamic_cast<SimpleLoopController&>(manager.loopController());
 714
 715   auto loopFunc = [&]() {
 716     if (!taskAdded) {
 717       manager.addTask([&]() {
 718         std::vector<std::function<void()>> funcs;
 719         for (size_t i = 0; i < 3; ++i) {
 720           funcs.push_back([&pendingFibers]() {
 721             await([&pendingFibers](Promise<int> promise) {
 722               pendingFibers.push_back(std::move(promise));
 723             });
 724           });
 725         }
 726
 727         auto results = collectN(funcs.begin(), funcs.end(), 2);
 728         EXPECT_EQ(2, results.size());
 729         EXPECT_EQ(1, pendingFibers.size());
 730       });
 731       taskAdded = true;
 732     } else if (pendingFibers.size()) {
 733       pendingFibers.back().setValue(0);
 734       pendingFibers.pop_back();
 735     } else {
 736       loopController.stop();
 737     }
 738   };
 739
 740   loopController.loop(std::move(loopFunc));
 741 }
 742
 743 TEST(FiberManager, collectNVoidThrow) {
 744   std::vector<Promise<int>> pendingFibers;
 745   bool taskAdded = false;
 746
 747   FiberManager manager(std::make_unique<SimpleLoopController>());
 748   auto& loopController =
 749       dynamic_cast<SimpleLoopController&>(manager.loopController());
 750
 751   auto loopFunc = [&]() {
 752     if (!taskAdded) {
 753       manager.addTask([&]() {
 754         std::vector<std::function<void()>> funcs;
 755         for (size_t i = 0; i < 3; ++i) {
 756           funcs.push_back([&pendingFibers]() {
 757             await([&pendingFibers](Promise<int> promise) {
 758               pendingFibers.push_back(std::move(promise));
 759             });
 760             throw std::runtime_error("Runtime");
 761           });
 762         }
 763
 764         try {
 765           collectN(funcs.begin(), funcs.end(), 2);
 766         } catch (...) {
 767           EXPECT_EQ(1, pendingFibers.size());
 768         }
 769       });
 770       taskAdded = true;
 771     } else if (pendingFibers.size()) {
 772       pendingFibers.back().setValue(0);
 773       pendingFibers.pop_back();
 774     } else {
 775       loopController.stop();
 776     }
 777   };
 778
 779   loopController.loop(std::move(loopFunc));
 780 }
 781
 782 TEST(FiberManager, collectAll) {
 783   std::vector<Promise<int>> pendingFibers;
 784   bool taskAdded = false;
 785
 786   FiberManager manager(std::make_unique<SimpleLoopController>());
 787   auto& loopController =
 788       dynamic_cast<SimpleLoopController&>(manager.loopController());
 789
 790   auto loopFunc = [&]() {
 791     if (!taskAdded) {
 792       manager.addTask([&]() {
 793         std::vector<std::function<int()>> funcs;
 794         for (size_t i = 0; i < 3; ++i) {
 795           funcs.push_back([i, &pendingFibers]() {
 796             await([&pendingFibers](Promise<int> promise) {
 797               pendingFibers.push_back(std::move(promise));
 798             });
 799             return i * 2 + 1;
 800           });
 801         }
 802
 803         auto results = collectAll(funcs.begin(), funcs.end());
 804         EXPECT_TRUE(pendingFibers.empty());
 805         for (size_t i = 0; i < 3; ++i) {
 806           EXPECT_EQ(i * 2 + 1, results[i]);
 807         }
 808       });
 809       taskAdded = true;
 810     } else if (pendingFibers.size()) {
 811       pendingFibers.back().setValue(0);
 812       pendingFibers.pop_back();
 813     } else {
 814       loopController.stop();
 815     }
 816   };
 817
 818   loopController.loop(std::move(loopFunc));
 819 }
 820
 821 TEST(FiberManager, collectAllVoid) {
 822   std::vector<Promise<int>> pendingFibers;
 823   bool taskAdded = false;
 824
 825   FiberManager manager(std::make_unique<SimpleLoopController>());
 826   auto& loopController =
 827       dynamic_cast<SimpleLoopController&>(manager.loopController());
 828
 829   auto loopFunc = [&]() {
 830     if (!taskAdded) {
 831       manager.addTask([&]() {
 832         std::vector<std::function<void()>> funcs;
 833         for (size_t i = 0; i < 3; ++i) {
 834           funcs.push_back([&pendingFibers]() {
 835             await([&pendingFibers](Promise<int> promise) {
 836               pendingFibers.push_back(std::move(promise));
 837             });
 838           });
 839         }
 840
 841         collectAll(funcs.begin(), funcs.end());
 842         EXPECT_TRUE(pendingFibers.empty());
 843       });
 844       taskAdded = true;
 845     } else if (pendingFibers.size()) {
 846       pendingFibers.back().setValue(0);
 847       pendingFibers.pop_back();
 848     } else {
 849       loopController.stop();
 850     }
 851   };
 852
 853   loopController.loop(std::move(loopFunc));
 854 }
 855
 856 TEST(FiberManager, collectAny) {
 857   std::vector<Promise<int>> pendingFibers;
 858   bool taskAdded = false;
 859
 860   FiberManager manager(std::make_unique<SimpleLoopController>());
 861   auto& loopController =
 862       dynamic_cast<SimpleLoopController&>(manager.loopController());
 863
 864   auto loopFunc = [&]() {
 865     if (!taskAdded) {
 866       manager.addTask([&]() {
 867         std::vector<std::function<int()>> funcs;
 868         for (size_t i = 0; i < 3; ++i) {
 869           funcs.push_back([i, &pendingFibers]() {
 870             await([&pendingFibers](Promise<int> promise) {
 871               pendingFibers.push_back(std::move(promise));
 872             });
 873             if (i == 1) {
 874               throw std::runtime_error("This exception will be ignored");
 875             }
 876             return i * 2 + 1;
 877           });
 878         }
 879
 880         auto result = collectAny(funcs.begin(), funcs.end());
 881         EXPECT_EQ(2, pendingFibers.size());
 882         EXPECT_EQ(2, result.first);
 883         EXPECT_EQ(2 * 2 + 1, result.second);
 884       });
 885       taskAdded = true;
 886     } else if (pendingFibers.size()) {
 887       pendingFibers.back().setValue(0);
 888       pendingFibers.pop_back();
 889     } else {
 890       loopController.stop();
 891     }
 892   };
 893
 894   loopController.loop(std::move(loopFunc));
 895 }
 896
 897 namespace {
 898 /* Checks that this function was run from a main context,
 899    by comparing an address on a stack to a known main stack address
 900    and a known related fiber stack address.  The assumption
 901    is that fiber stack and main stack will be far enough apart,
 902    while any two values on the same stack will be close. */
 903 void expectMainContext(bool& ran, int* mainLocation, int* fiberLocation) {
 904   int here;
 905   /* 2 pages is a good guess */
 906   constexpr auto const kHereToFiberMaxDist = 0x2000 / sizeof(int);
 907
 908   // With ASAN's detect_stack_use_after_return=1, this must be much larger
 909   // I measured 410028 on x86_64, so allow for quadruple that, just in case.
 910   constexpr auto const kHereToMainMaxDist =
 911       folly::kIsSanitizeAddress ? 4 * 410028 : kHereToFiberMaxDist;
 912
 913   if (fiberLocation) {
 914     EXPECT_GT(std::abs(&here - fiberLocation), kHereToFiberMaxDist);
 915   }
 916   if (mainLocation) {
 917     EXPECT_LT(std::abs(&here - mainLocation), kHereToMainMaxDist);
 918   }
 919
 920   EXPECT_FALSE(ran);
 921   ran = true;
 922 }
 923 } // namespace
 924
 925 TEST(FiberManager, runInMainContext) {
 926   FiberManager manager(std::make_unique<SimpleLoopController>());
 927   auto& loopController =
 928       dynamic_cast<SimpleLoopController&>(manager.loopController());
 929
 930   bool checkRan = false;
 931
 932   int mainLocation;
 933   manager.runInMainContext(
 934       [&]() { expectMainContext(checkRan, &mainLocation, nullptr); });
 935   EXPECT_TRUE(checkRan);
 936
 937   checkRan = false;
 938
 939   manager.addTask([&]() {
 940     struct A {
 941       explicit A(int value_) : value(value_) {}
 942       A(const A&) = delete;
 943       A(A&&) = default;
 944
 945       int value;
 946     };
 947     int stackLocation;
 948     auto ret = runInMainContext([&]() {
 949       expectMainContext(checkRan, &mainLocation, &stackLocation);
 950       return A(42);
 951     });
 952     EXPECT_TRUE(checkRan);
 953     EXPECT_EQ(42, ret.value);
 954   });
 955
 956   loopController.loop([&]() { loopController.stop(); });
 957
 958   EXPECT_TRUE(checkRan);
 959 }
 960
 961 TEST(FiberManager, addTaskFinally) {
 962   FiberManager manager(std::make_unique<SimpleLoopController>());
 963   auto& loopController =
 964       dynamic_cast<SimpleLoopController&>(manager.loopController());
 965
 966   bool checkRan = false;
 967
 968   int mainLocation;
 969
 970   manager.addTaskFinally(
 971       [&]() { return 1234; },
 972       [&](Try<int>&& result) {
 973         EXPECT_EQ(result.value(), 1234);
 974
 975         expectMainContext(checkRan, &mainLocation, nullptr);
 976       });
 977
 978   EXPECT_FALSE(checkRan);
 979
 980   loopController.loop([&]() { loopController.stop(); });
 981
 982   EXPECT_TRUE(checkRan);
 983 }
 984
 985 TEST(FiberManager, fibersPoolWithinLimit) {
 986   FiberManager::Options opts;
 987   opts.maxFibersPoolSize = 5;
 988
 989   FiberManager manager(std::make_unique<SimpleLoopController>(), opts);
 990   auto& loopController =
 991       dynamic_cast<SimpleLoopController&>(manager.loopController());
 992
 993   size_t fibersRun = 0;
 994
 995   for (size_t i = 0; i < 5; ++i) {
 996     manager.addTask([&]() { ++fibersRun; });
 997   }
 998   loopController.loop([&]() { loopController.stop(); });
 999
1000   EXPECT_EQ(5, fibersRun);
1001   EXPECT_EQ(5, manager.fibersAllocated());
1002   EXPECT_EQ(5, manager.fibersPoolSize());
1003
1004   for (size_t i = 0; i < 5; ++i) {
1005     manager.addTask([&]() { ++fibersRun; });
1006   }
1007   loopController.loop([&]() { loopController.stop(); });
1008
1009   EXPECT_EQ(10, fibersRun);
1010   EXPECT_EQ(5, manager.fibersAllocated());
1011   EXPECT_EQ(5, manager.fibersPoolSize());
1012 }
1013
1014 TEST(FiberManager, fibersPoolOverLimit) {
1015   FiberManager::Options opts;
1016   opts.maxFibersPoolSize = 5;
1017
1018   FiberManager manager(std::make_unique<SimpleLoopController>(), opts);
1019   auto& loopController =
1020       dynamic_cast<SimpleLoopController&>(manager.loopController());
1021
1022   size_t fibersRun = 0;
1023
1024   for (size_t i = 0; i < 10; ++i) {
1025     manager.addTask([&]() { ++fibersRun; });
1026   }
1027
1028   EXPECT_EQ(0, fibersRun);
1029   EXPECT_EQ(10, manager.fibersAllocated());
1030   EXPECT_EQ(0, manager.fibersPoolSize());
1031
1032   loopController.loop([&]() { loopController.stop(); });
1033
1034   EXPECT_EQ(10, fibersRun);
1035   EXPECT_EQ(5, manager.fibersAllocated());
1036   EXPECT_EQ(5, manager.fibersPoolSize());
1037 }
1038
1039 TEST(FiberManager, remoteFiberBasic) {
1040   FiberManager manager(std::make_unique<SimpleLoopController>());
1041   auto& loopController =
1042       dynamic_cast<SimpleLoopController&>(manager.loopController());
1043
1044   int result[2];
1045   result[0] = result[1] = 0;
1046   folly::Optional<Promise<int>> savedPromise[2];
1047   manager.addTask([&]() {
1048     result[0] = await(
1049         [&](Promise<int> promise) { savedPromise[0] = std::move(promise); });
1050   });
1051   manager.addTask([&]() {
1052     result[1] = await(
1053         [&](Promise<int> promise) { savedPromise[1] = std::move(promise); });
1054   });
1055
1056   manager.loopUntilNoReady();
1057
1058   EXPECT_TRUE(savedPromise[0].hasValue());
1059   EXPECT_TRUE(savedPromise[1].hasValue());
1060   EXPECT_EQ(0, result[0]);
1061   EXPECT_EQ(0, result[1]);
1062
1063   std::thread remoteThread0{[&]() { savedPromise[0]->setValue(42); }};
1064   std::thread remoteThread1{[&]() { savedPromise[1]->setValue(43); }};
1065   remoteThread0.join();
1066   remoteThread1.join();
1067   EXPECT_EQ(0, result[0]);
1068   EXPECT_EQ(0, result[1]);
1069   /* Should only have scheduled once */
1070   EXPECT_EQ(1, loopController.remoteScheduleCalled());
1071
1072   manager.loopUntilNoReady();
1073   EXPECT_EQ(42, result[0]);
1074   EXPECT_EQ(43, result[1]);
1075 }
1076
1077 TEST(FiberManager, addTaskRemoteBasic) {
1078   FiberManager manager(std::make_unique<SimpleLoopController>());
1079
1080   int result[2];
1081   result[0] = result[1] = 0;
1082   folly::Optional<Promise<int>> savedPromise[2];
1083
1084   std::thread remoteThread0{[&]() {
1085     manager.addTaskRemote([&]() {
1086       result[0] = await(
1087           [&](Promise<int> promise) { savedPromise[0] = std::move(promise); });
1088     });
1089   }};
1090   std::thread remoteThread1{[&]() {
1091     manager.addTaskRemote([&]() {
1092       result[1] = await(
1093           [&](Promise<int> promise) { savedPromise[1] = std::move(promise); });
1094     });
1095   }};
1096   remoteThread0.join();
1097   remoteThread1.join();
1098
1099   manager.loopUntilNoReady();
1100
1101   EXPECT_TRUE(savedPromise[0].hasValue());
1102   EXPECT_TRUE(savedPromise[1].hasValue());
1103   EXPECT_EQ(0, result[0]);
1104   EXPECT_EQ(0, result[1]);
1105
1106   savedPromise[0]->setValue(42);
1107   savedPromise[1]->setValue(43);
1108
1109   EXPECT_EQ(0, result[0]);
1110   EXPECT_EQ(0, result[1]);
1111
1112   manager.loopUntilNoReady();
1113   EXPECT_EQ(42, result[0]);
1114   EXPECT_EQ(43, result[1]);
1115 }
1116
1117 TEST(FiberManager, remoteHasTasks) {
1118   size_t counter = 0;
1119   FiberManager fm(std::make_unique<SimpleLoopController>());
1120   std::thread remote([&]() { fm.addTaskRemote([&]() { ++counter; }); });
1121
1122   remote.join();
1123
1124   while (fm.hasTasks()) {
1125     fm.loopUntilNoReady();
1126   }
1127
1128   EXPECT_FALSE(fm.hasTasks());
1129   EXPECT_EQ(counter, 1);
1130 }
1131
1132 TEST(FiberManager, remoteHasReadyTasks) {
1133   int result = 0;
1134   folly::Optional<Promise<int>> savedPromise;
1135   FiberManager fm(std::make_unique<SimpleLoopController>());
1136   std::thread remote([&]() {
1137     fm.addTaskRemote([&]() {
1138       result = await(
1139           [&](Promise<int> promise) { savedPromise = std::move(promise); });
1140       EXPECT_TRUE(fm.hasTasks());
1141     });
1142   });
1143
1144   remote.join();
1145   EXPECT_TRUE(fm.hasTasks());
1146
1147   fm.loopUntilNoReady();
1148   EXPECT_TRUE(fm.hasTasks());
1149
1150   std::thread remote2([&]() { savedPromise->setValue(47); });
1151   remote2.join();
1152   EXPECT_TRUE(fm.hasTasks());
1153
1154   fm.loopUntilNoReady();
1155   EXPECT_FALSE(fm.hasTasks());
1156
1157   EXPECT_EQ(result, 47);
1158 }
1159
1160 template <typename Data>
1161 void testFiberLocal() {
1162   FiberManager fm(LocalType<Data>(), std::make_unique<SimpleLoopController>());
1163
1164   fm.addTask([]() {
1165     EXPECT_EQ(42, local<Data>().value);
1166
1167     local<Data>().value = 43;
1168
1169     addTask([]() {
1170       EXPECT_EQ(43, local<Data>().value);
1171
1172       local<Data>().value = 44;
1173
1174       addTask([]() { EXPECT_EQ(44, local<Data>().value); });
1175     });
1176   });
1177
1178   fm.addTask([&]() {
1179     EXPECT_EQ(42, local<Data>().value);
1180
1181     local<Data>().value = 43;
1182
1183     fm.addTaskRemote([]() { EXPECT_EQ(43, local<Data>().value); });
1184   });
1185
1186   fm.addTask([]() {
1187     EXPECT_EQ(42, local<Data>().value);
1188     local<Data>().value = 43;
1189
1190     auto task = []() {
1191       EXPECT_EQ(43, local<Data>().value);
1192       local<Data>().value = 44;
1193     };
1194     std::vector<std::function<void()>> tasks{task};
1195     collectAny(tasks.begin(), tasks.end());
1196
1197     EXPECT_EQ(43, local<Data>().value);
1198   });
1199
1200   fm.loopUntilNoReady();
1201   EXPECT_FALSE(fm.hasTasks());
1202 }
1203
1204 TEST(FiberManager, fiberLocal) {
1205   struct SimpleData {
1206     int value{42};
1207   };
1208
1209   testFiberLocal<SimpleData>();
1210 }
1211
1212 TEST(FiberManager, fiberLocalHeap) {
1213   struct LargeData {
1214     char _[1024 * 1024];
1215     int value{42};
1216   };
1217
1218   testFiberLocal<LargeData>();
1219 }
1220
1221 TEST(FiberManager, fiberLocalDestructor) {
1222   struct CrazyData {
1223     size_t data{42};
1224
1225     ~CrazyData() {
1226       if (data == 41) {
1227         addTask([]() {
1228           EXPECT_EQ(42, local<CrazyData>().data);
1229           // Make sure we don't have infinite loop
1230           local<CrazyData>().data = 0;
1231         });
1232       }
1233     }
1234   };
1235
1236   FiberManager fm(
1237       LocalType<CrazyData>(), std::make_unique<SimpleLoopController>());
1238
1239   fm.addTask([]() { local<CrazyData>().data = 41; });
1240
1241   fm.loopUntilNoReady();
1242   EXPECT_FALSE(fm.hasTasks());
1243 }
1244
1245 TEST(FiberManager, yieldTest) {
1246   FiberManager manager(std::make_unique<SimpleLoopController>());
1247   auto& loopController =
1248       dynamic_cast<SimpleLoopController&>(manager.loopController());
1249
1250   bool checkRan = false;
1251
1252   manager.addTask([&]() {
1253     manager.yield();
1254     checkRan = true;
1255   });
1256
1257   loopController.loop([&]() {
1258     if (checkRan) {
1259       loopController.stop();
1260     }
1261   });
1262
1263   EXPECT_TRUE(checkRan);
1264 }
1265
1266 TEST(FiberManager, RequestContext) {
1267   FiberManager fm(std::make_unique<SimpleLoopController>());
1268
1269   bool checkRun1 = false;
1270   bool checkRun2 = false;
1271   bool checkRun3 = false;
1272   bool checkRun4 = false;
1273   folly::fibers::Baton baton1;
1274   folly::fibers::Baton baton2;
1275   folly::fibers::Baton baton3;
1276   folly::fibers::Baton baton4;
1277
1278   {
1279     folly::RequestContextScopeGuard rctx;
1280     auto rcontext1 = folly::RequestContext::get();
1281     fm.addTask([&, rcontext1]() {
1282       EXPECT_EQ(rcontext1, folly::RequestContext::get());
1283       baton1.wait(
1284           [&]() { EXPECT_EQ(rcontext1, folly::RequestContext::get()); });
1285       EXPECT_EQ(rcontext1, folly::RequestContext::get());
1286       runInMainContext(
1287           [&]() { EXPECT_EQ(rcontext1, folly::RequestContext::get()); });
1288       checkRun1 = true;
1289     });
1290   }
1291   {
1292     folly::RequestContextScopeGuard rctx;
1293     auto rcontext2 = folly::RequestContext::get();
1294     fm.addTaskRemote([&, rcontext2]() {
1295       EXPECT_EQ(rcontext2, folly::RequestContext::get());
1296       baton2.wait();
1297       EXPECT_EQ(rcontext2, folly::RequestContext::get());
1298       checkRun2 = true;
1299     });
1300   }
1301   {
1302     folly::RequestContextScopeGuard rctx;
1303     auto rcontext3 = folly::RequestContext::get();
1304     fm.addTaskFinally(
1305         [&, rcontext3]() {
1306           EXPECT_EQ(rcontext3, folly::RequestContext::get());
1307           baton3.wait();
1308           EXPECT_EQ(rcontext3, folly::RequestContext::get());
1309
1310           return folly::Unit();
1311         },
1312         [&, rcontext3](Try<folly::Unit>&& /* t */) {
1313           EXPECT_EQ(rcontext3, folly::RequestContext::get());
1314           checkRun3 = true;
1315         });
1316   }
1317   {
1318     folly::RequestContext::setContext(nullptr);
1319     fm.addTask([&]() {
1320       folly::RequestContextScopeGuard rctx;
1321       auto rcontext4 = folly::RequestContext::get();
1322       baton4.wait();
1323       EXPECT_EQ(rcontext4, folly::RequestContext::get());
1324       checkRun4 = true;
1325     });
1326   }
1327   {
1328     folly::RequestContextScopeGuard rctx;
1329     auto rcontext = folly::RequestContext::get();
1330
1331     fm.loopUntilNoReady();
1332     EXPECT_EQ(rcontext, folly::RequestContext::get());
1333
1334     baton1.post();
1335     EXPECT_EQ(rcontext, folly::RequestContext::get());
1336     fm.loopUntilNoReady();
1337     EXPECT_TRUE(checkRun1);
1338     EXPECT_EQ(rcontext, folly::RequestContext::get());
1339
1340     baton2.post();
1341     EXPECT_EQ(rcontext, folly::RequestContext::get());
1342     fm.loopUntilNoReady();
1343     EXPECT_TRUE(checkRun2);
1344     EXPECT_EQ(rcontext, folly::RequestContext::get());
1345
1346     baton3.post();
1347     EXPECT_EQ(rcontext, folly::RequestContext::get());
1348     fm.loopUntilNoReady();
1349     EXPECT_TRUE(checkRun3);
1350     EXPECT_EQ(rcontext, folly::RequestContext::get());
1351
1352     baton4.post();
1353     EXPECT_EQ(rcontext, folly::RequestContext::get());
1354     fm.loopUntilNoReady();
1355     EXPECT_TRUE(checkRun4);
1356     EXPECT_EQ(rcontext, folly::RequestContext::get());
1357   }
1358 }
1359
1360 TEST(FiberManager, resizePeriodically) {
1361   FiberManager::Options opts;
1362   opts.fibersPoolResizePeriodMs = 300;
1363   opts.maxFibersPoolSize = 5;
1364
1365   FiberManager manager(std::make_unique<EventBaseLoopController>(), opts);
1366
1367   folly::EventBase evb;
1368   dynamic_cast<EventBaseLoopController&>(manager.loopController())
1369       .attachEventBase(evb);
1370
1371   std::vector<Baton> batons(10);
1372
1373   size_t tasksRun = 0;
1374   for (size_t i = 0; i < 30; ++i) {
1375     manager.addTask([i, &batons, &tasksRun]() {
1376       ++tasksRun;
1377       // Keep some fibers active indefinitely
1378       if (i < batons.size()) {
1379         batons[i].wait();
1380       }
1381     });
1382   }
1383
1384   EXPECT_EQ(0, tasksRun);
1385   EXPECT_EQ(30, manager.fibersAllocated());
1386   EXPECT_EQ(0, manager.fibersPoolSize());
1387
1388   evb.loopOnce();
1389   EXPECT_EQ(30, tasksRun);
1390   EXPECT_EQ(30, manager.fibersAllocated());
1391   // Can go over maxFibersPoolSize, 10 of 30 fibers still active
1392   EXPECT_EQ(20, manager.fibersPoolSize());
1393
1394   std::this_thread::sleep_for(std::chrono::milliseconds(400));
1395   evb.loopOnce(); // no fibers active in this period
1396   EXPECT_EQ(30, manager.fibersAllocated());
1397   EXPECT_EQ(20, manager.fibersPoolSize());
1398
1399   std::this_thread::sleep_for(std::chrono::milliseconds(400));
1400   evb.loopOnce(); // should shrink fibers pool to maxFibersPoolSize
1401   EXPECT_EQ(15, manager.fibersAllocated());
1402   EXPECT_EQ(5, manager.fibersPoolSize());
1403
1404   for (size_t i = 0; i < batons.size(); ++i) {
1405     batons[i].post();
1406   }
1407   evb.loopOnce();
1408   EXPECT_EQ(15, manager.fibersAllocated());
1409   EXPECT_EQ(15, manager.fibersPoolSize());
1410
1411   std::this_thread::sleep_for(std::chrono::milliseconds(400));
1412   evb.loopOnce(); // 10 fibers active in last period
1413   EXPECT_EQ(10, manager.fibersAllocated());
1414   EXPECT_EQ(10, manager.fibersPoolSize());
1415
1416   std::this_thread::sleep_for(std::chrono::milliseconds(400));
1417   evb.loopOnce();
1418   EXPECT_EQ(5, manager.fibersAllocated());
1419   EXPECT_EQ(5, manager.fibersPoolSize());
1420 }
1421
1422 TEST(FiberManager, batonWaitTimeoutHandler) {
1423   FiberManager manager(std::make_unique<EventBaseLoopController>());
1424
1425   folly::EventBase evb;
1426   dynamic_cast<EventBaseLoopController&>(manager.loopController())
1427       .attachEventBase(evb);
1428
1429   size_t fibersRun = 0;
1430   Baton baton;
1431   Baton::TimeoutHandler timeoutHandler;
1432
1433   manager.addTask([&]() {
1434     baton.wait(timeoutHandler);
1435     ++fibersRun;
1436   });
1437   manager.loopUntilNoReady();
1438
1439   EXPECT_FALSE(baton.try_wait());
1440   EXPECT_EQ(0, fibersRun);
1441
1442   timeoutHandler.scheduleTimeout(std::chrono::milliseconds(250));
1443   std::this_thread::sleep_for(std::chrono::milliseconds(500));
1444
1445   EXPECT_FALSE(baton.try_wait());
1446   EXPECT_EQ(0, fibersRun);
1447
1448   evb.loopOnce();
1449   manager.loopUntilNoReady();
1450
1451   EXPECT_EQ(1, fibersRun);
1452 }
1453
1454 TEST(FiberManager, batonWaitTimeoutMany) {
1455   FiberManager manager(std::make_unique<EventBaseLoopController>());
1456
1457   folly::EventBase evb;
1458   dynamic_cast<EventBaseLoopController&>(manager.loopController())
1459       .attachEventBase(evb);
1460
1461   constexpr size_t kNumTimeoutTasks = 10000;
1462   size_t tasksCount = kNumTimeoutTasks;
1463
1464   // We add many tasks to hit timeout queue deallocation logic.
1465   for (size_t i = 0; i < kNumTimeoutTasks; ++i) {
1466     manager.addTask([&]() {
1467       Baton baton;
1468       Baton::TimeoutHandler timeoutHandler;
1469
1470       folly::fibers::addTask([&] {
1471         timeoutHandler.scheduleTimeout(std::chrono::milliseconds(1000));
1472       });
1473
1474       baton.wait(timeoutHandler);
1475       if (--tasksCount == 0) {
1476         evb.terminateLoopSoon();
1477       }
1478     });
1479   }
1480
1481   evb.loopForever();
1482 }
1483
1484 TEST(FiberManager, remoteFutureTest) {
1485   FiberManager fiberManager(std::make_unique<SimpleLoopController>());
1486   auto& loopController =
1487       dynamic_cast<SimpleLoopController&>(fiberManager.loopController());
1488
1489   int testValue1 = 5;
1490   int testValue2 = 7;
1491   auto f1 = fiberManager.addTaskFuture([&]() { return testValue1; });
1492   auto f2 = fiberManager.addTaskRemoteFuture([&]() { return testValue2; });
1493   loopController.loop([&]() { loopController.stop(); });
1494   auto v1 = f1.get();
1495   auto v2 = f2.get();
1496
1497   EXPECT_EQ(v1, testValue1);
1498   EXPECT_EQ(v2, testValue2);
1499 }
1500
1501 // Test that a void function produes a Future<Unit>.
1502 TEST(FiberManager, remoteFutureVoidUnitTest) {
1503   FiberManager fiberManager(std::make_unique<SimpleLoopController>());
1504   auto& loopController =
1505       dynamic_cast<SimpleLoopController&>(fiberManager.loopController());
1506
1507   bool ranLocal = false;
1508   folly::Future<folly::Unit> futureLocal =
1509       fiberManager.addTaskFuture([&]() { ranLocal = true; });
1510
1511   bool ranRemote = false;
1512   folly::Future<folly::Unit> futureRemote =
1513       fiberManager.addTaskRemoteFuture([&]() { ranRemote = true; });
1514
1515   loopController.loop([&]() { loopController.stop(); });
1516
1517   futureLocal.wait();
1518   ASSERT_TRUE(ranLocal);
1519
1520   futureRemote.wait();
1521   ASSERT_TRUE(ranRemote);
1522 }
1523
1524 TEST(FiberManager, nestedFiberManagers) {
1525   folly::EventBase outerEvb;
1526   folly::EventBase innerEvb;
1527
1528   getFiberManager(outerEvb).addTask([&]() {
1529     EXPECT_EQ(
1530         &getFiberManager(outerEvb), FiberManager::getFiberManagerUnsafe());
1531
1532     runInMainContext([&]() {
1533       getFiberManager(innerEvb).addTask([&]() {
1534         EXPECT_EQ(
1535             &getFiberManager(innerEvb), FiberManager::getFiberManagerUnsafe());
1536
1537         innerEvb.terminateLoopSoon();
1538       });
1539
1540       innerEvb.loopForever();
1541     });
1542
1543     EXPECT_EQ(
1544         &getFiberManager(outerEvb), FiberManager::getFiberManagerUnsafe());
1545
1546     outerEvb.terminateLoopSoon();
1547   });
1548
1549   outerEvb.loopForever();
1550 }
1551
1552 TEST(FiberManager, semaphore) {
1553   static constexpr size_t kTasks = 10;
1554   static constexpr size_t kIterations = 10000;
1555   static constexpr size_t kNumTokens = 10;
1556
1557   Semaphore sem(kNumTokens);
1558   int counterA = 0;
1559   int counterB = 0;
1560
1561   auto task = [&sem](int& counter, folly::fibers::Baton& baton) {
1562     FiberManager manager(std::make_unique<EventBaseLoopController>());
1563     folly::EventBase evb;
1564     dynamic_cast<EventBaseLoopController&>(manager.loopController())
1565         .attachEventBase(evb);
1566
1567     {
1568       std::shared_ptr<folly::EventBase> completionCounter(
1569           &evb, [](folly::EventBase* evb) { evb->terminateLoopSoon(); });
1570
1571       for (size_t i = 0; i < kTasks; ++i) {
1572         manager.addTask([&, completionCounter]() {
1573           for (size_t j = 0; j < kIterations; ++j) {
1574             sem.wait();
1575             ++counter;
1576             sem.signal();
1577             --counter;
1578
1579             EXPECT_LT(counter, kNumTokens);
1580             EXPECT_GE(counter, 0);
1581           }
1582         });
1583       }
1584
1585       baton.wait();
1586     }
1587     evb.loopForever();
1588   };
1589
1590   folly::fibers::Baton batonA;
1591   folly::fibers::Baton batonB;
1592   std::thread threadA([&] { task(counterA, batonA); });
1593   std::thread threadB([&] { task(counterB, batonB); });
1594
1595   batonA.post();
1596   batonB.post();
1597   threadA.join();
1598   threadB.join();
1599
1600   EXPECT_LT(counterA, kNumTokens);
1601   EXPECT_LT(counterB, kNumTokens);
1602   EXPECT_GE(counterA, 0);
1603   EXPECT_GE(counterB, 0);
1604 }
1605
1606 template <typename ExecutorT>
1607 void singleBatchDispatch(ExecutorT& executor, int batchSize, int index) {
1608   thread_local BatchDispatcher<int, std::string, ExecutorT> batchDispatcher(
1609       executor, [=](std::vector<int>&& batch) {
1610         EXPECT_EQ(batchSize, batch.size());
1611         std::vector<std::string> results;
1612         for (auto& it : batch) {
1613           results.push_back(folly::to<std::string>(it));
1614         }
1615         return results;
1616       });
1617
1618   auto indexCopy = index;
1619   auto result = batchDispatcher.add(std::move(indexCopy));
1620   EXPECT_EQ(folly::to<std::string>(index), result.get());
1621 }
1622
1623 TEST(FiberManager, batchDispatchTest) {
1624   folly::EventBase evb;
1625   auto& executor = getFiberManager(evb);
1626
1627   // Launch multiple fibers with a single id.
1628   executor.add([&]() {
1629     int batchSize = 10;
1630     for (int i = 0; i < batchSize; i++) {
1631       executor.add(
1632           [=, &executor]() { singleBatchDispatch(executor, batchSize, i); });
1633     }
1634   });
1635   evb.loop();
1636
1637   // Reuse the same BatchDispatcher to batch once again.
1638   executor.add([&]() {
1639     int batchSize = 10;
1640     for (int i = 0; i < batchSize; i++) {
1641       executor.add(
1642           [=, &executor]() { singleBatchDispatch(executor, batchSize, i); });
1643     }
1644   });
1645   evb.loop();
1646 }
1647
1648 template <typename ExecutorT>
1649 folly::Future<std::vector<std::string>> doubleBatchInnerDispatch(
1650     ExecutorT& executor,
1651     int totalNumberOfElements,
1652     std::vector<int> input) {
1653   thread_local BatchDispatcher<
1654       std::vector<int>,
1655       std::vector<std::string>,
1656       ExecutorT>
1657       batchDispatcher(executor, [=](std::vector<std::vector<int>>&& batch) {
1658         std::vector<std::vector<std::string>> results;
1659         int numberOfElements = 0;
1660         for (auto& unit : batch) {
1661           numberOfElements += unit.size();
1662           std::vector<std::string> result;
1663           for (auto& element : unit) {
1664             result.push_back(folly::to<std::string>(element));
1665           }
1666           results.push_back(std::move(result));
1667         }
1668         EXPECT_EQ(totalNumberOfElements, numberOfElements);
1669         return results;
1670       });
1671
1672   return batchDispatcher.add(std::move(input));
1673 }
1674
1675 /**
1676  * Batch values in groups of 5, and then call inner dispatch.
1677  */
1678 template <typename ExecutorT>
1679 void doubleBatchOuterDispatch(
1680     ExecutorT& executor,
1681     int totalNumberOfElements,
1682     int index) {
1683   thread_local BatchDispatcher<int, std::string, ExecutorT> batchDispatcher(
1684       executor, [=, &executor](std::vector<int>&& batch) {
1685         EXPECT_EQ(totalNumberOfElements, batch.size());
1686         std::vector<std::string> results;
1687         std::vector<folly::Future<std::vector<std::string>>>
1688             innerDispatchResultFutures;
1689
1690         std::vector<int> group;
1691         for (auto unit : batch) {
1692           group.push_back(unit);
1693           if (group.size() == 5) {
1694             auto localGroup = group;
1695             group.clear();
1696
1697             innerDispatchResultFutures.push_back(doubleBatchInnerDispatch(
1698                 executor, totalNumberOfElements, localGroup));
1699           }
1700         }
1701
1702         folly::collectAll(
1703             innerDispatchResultFutures.begin(),
1704             innerDispatchResultFutures.end())
1705             .then([&](std::vector<Try<std::vector<std::string>>>
1706                           innerDispatchResults) {
1707               for (auto& unit : innerDispatchResults) {
1708                 for (auto& element : unit.value()) {
1709                   results.push_back(element);
1710                 }
1711               }
1712             })
1713             .get();
1714         return results;
1715       });
1716
1717   auto indexCopy = index;
1718   auto result = batchDispatcher.add(std::move(indexCopy));
1719   EXPECT_EQ(folly::to<std::string>(index), result.get());
1720 }
1721
1722 TEST(FiberManager, doubleBatchDispatchTest) {
1723   folly::EventBase evb;
1724   auto& executor = getFiberManager(evb);
1725
1726   // Launch multiple fibers with a single id.
1727   executor.add([&]() {
1728     int totalNumberOfElements = 20;
1729     for (int i = 0; i < totalNumberOfElements; i++) {
1730       executor.add([=, &executor]() {
1731         doubleBatchOuterDispatch(executor, totalNumberOfElements, i);
1732       });
1733     }
1734   });
1735   evb.loop();
1736 }
1737
1738 template <typename ExecutorT>
1739 void batchDispatchExceptionHandling(ExecutorT& executor, int i) {
1740   thread_local BatchDispatcher<int, int, ExecutorT> batchDispatcher(
1741       executor, [](std::vector<int> &&) -> std::vector<int> {
1742         throw std::runtime_error("Surprise!!");
1743       });
1744
1745   EXPECT_THROW(batchDispatcher.add(i).get(), std::runtime_error);
1746 }
1747
1748 TEST(FiberManager, batchDispatchExceptionHandlingTest) {
1749   folly::EventBase evb;
1750   auto& executor = getFiberManager(evb);
1751
1752   // Launch multiple fibers with a single id.
1753   executor.add([&]() {
1754     int totalNumberOfElements = 5;
1755     for (int i = 0; i < totalNumberOfElements; i++) {
1756       executor.add(
1757           [=, &executor]() { batchDispatchExceptionHandling(executor, i); });
1758     }
1759   });
1760   evb.loop();
1761 }
1762
1763 namespace AtomicBatchDispatcherTesting {
1764
1765 using ValueT = size_t;
1766 using ResultT = std::string;
1767 using DispatchFunctionT =
1768     folly::Function<std::vector<ResultT>(std::vector<ValueT>&&)>;
1769
1770 #define ENABLE_TRACE_IN_TEST 0 // Set to 1 to debug issues in ABD tests
1771 #if ENABLE_TRACE_IN_TEST
1772 #define OUTPUT_TRACE std::cerr
1773 #else // ENABLE_TRACE_IN_TEST
1774 struct DevNullPiper {
1775   template <typename T>
1776   DevNullPiper& operator<<(const T&) {
1777     return *this;
1778   }
1779
1780   DevNullPiper& operator<<(std::ostream& (*)(std::ostream&)) {
1781     return *this;
1782   }
1783 } devNullPiper;
1784 #define OUTPUT_TRACE devNullPiper
1785 #endif // ENABLE_TRACE_IN_TEST
1786
1787 struct Job {
1788   AtomicBatchDispatcher<ValueT, ResultT>::Token token;
1789   ValueT input;
1790
1791   void preprocess(FiberManager& executor, bool die) {
1792     // Yield for a random duration [0, 10] ms to simulate I/O in preprocessing
1793     clock_t msecToDoIO = folly::Random::rand32() % 10;
1794     double start = (1000.0 * clock()) / CLOCKS_PER_SEC;
1795     double endAfter = start + msecToDoIO;
1796     while ((1000.0 * clock()) / CLOCKS_PER_SEC < endAfter) {
1797       executor.yield();
1798     }
1799     if (die) {
1800       throw std::logic_error("Simulating preprocessing failure");
1801     }
1802   }
1803
1804   Job(AtomicBatchDispatcher<ValueT, ResultT>::Token&& t, ValueT i)
1805       : token(std::move(t)), input(i) {}
1806
1807   Job(Job&&) = default;
1808   Job& operator=(Job&&) = default;
1809 };
1810
1811 ResultT processSingleInput(ValueT&& input) {
1812   return folly::to<ResultT>(std::move(input));
1813 }
1814
1815 std::vector<ResultT> userDispatchFunc(std::vector<ValueT>&& inputs) {
1816   size_t expectedCount = inputs.size();
1817   std::vector<ResultT> results;
1818   results.reserve(expectedCount);
1819   for (size_t i = 0; i < expectedCount; ++i) {
1820     results.emplace_back(processSingleInput(std::move(inputs[i])));
1821   }
1822   return results;
1823 }
1824
1825 void createJobs(
1826     AtomicBatchDispatcher<ValueT, ResultT>& atomicBatchDispatcher,
1827     std::vector<Job>& jobs,
1828     size_t count) {
1829   jobs.clear();
1830   for (size_t i = 0; i < count; ++i) {
1831     jobs.emplace_back(Job(atomicBatchDispatcher.getToken(), i));
1832   }
1833 }
1834
1835 enum class DispatchProblem {
1836   None,
1837   PreprocessThrows,
1838   DuplicateDispatch,
1839 };
1840
1841 void dispatchJobs(
1842     FiberManager& executor,
1843     std::vector<Job>& jobs,
1844     std::vector<folly::Optional<folly::Future<ResultT>>>& results,
1845     DispatchProblem dispatchProblem = DispatchProblem::None,
1846     size_t problemIndex = size_t(-1)) {
1847   EXPECT_TRUE(
1848       dispatchProblem == DispatchProblem::None || problemIndex < jobs.size());
1849   results.clear();
1850   results.resize(jobs.size());
1851   for (size_t i = 0; i < jobs.size(); ++i) {
1852     executor.add(
1853         [i, &executor, &jobs, &results, dispatchProblem, problemIndex]() {
1854           try {
1855             Job job(std::move(jobs[i]));
1856
1857             if (dispatchProblem == DispatchProblem::PreprocessThrows) {
1858               if (i == problemIndex) {
1859                 EXPECT_THROW(job.preprocess(executor, true), std::logic_error);
1860                 return;
1861               }
1862             }
1863
1864             job.preprocess(executor, false);
1865             OUTPUT_TRACE << "Dispatching job #" << i << std::endl;
1866             results[i] = job.token.dispatch(job.input);
1867             OUTPUT_TRACE << "Result future filled for job #" << i << std::endl;
1868
1869             if (dispatchProblem == DispatchProblem::DuplicateDispatch) {
1870               if (i == problemIndex) {
1871                 EXPECT_THROW(job.token.dispatch(job.input), ABDUsageException);
1872               }
1873             }
1874           } catch (...) {
1875             OUTPUT_TRACE << "Preprocessing failed for job #" << i << std::endl;
1876           }
1877         });
1878   }
1879 }
1880
1881 void validateResult(
1882     std::vector<folly::Optional<folly::Future<ResultT>>>& results,
1883     size_t i) {
1884   try {
1885     OUTPUT_TRACE << "results[" << i << "].value() : " << results[i]->value()
1886                  << std::endl;
1887   } catch (std::exception& e) {
1888     OUTPUT_TRACE << "Exception : " << e.what() << std::endl;
1889     throw;
1890   }
1891 }
1892
1893 template <typename TException>
1894 void validateResults(
1895     std::vector<folly::Optional<folly::Future<ResultT>>>& results,
1896     size_t expectedNumResults) {
1897   size_t numResultsFilled = 0;
1898   for (size_t i = 0; i < results.size(); ++i) {
1899     if (!results[i]) {
1900       continue;
1901     }
1902     ++numResultsFilled;
1903     EXPECT_THROW(validateResult(results, i), TException);
1904   }
1905   EXPECT_EQ(numResultsFilled, expectedNumResults);
1906 }
1907
1908 void validateResults(
1909     std::vector<folly::Optional<folly::Future<ResultT>>>& results,
1910     size_t expectedNumResults) {
1911   size_t numResultsFilled = 0;
1912   for (size_t i = 0; i < results.size(); ++i) {
1913     if (!results[i]) {
1914       continue;
1915     }
1916     ++numResultsFilled;
1917     EXPECT_NO_THROW(validateResult(results, i));
1918     ValueT expectedInput = i;
1919     EXPECT_EQ(
1920         results[i]->value(), processSingleInput(std::move(expectedInput)));
1921   }
1922   EXPECT_EQ(numResultsFilled, expectedNumResults);
1923 }
1924
1925 } // namespace AtomicBatchDispatcherTesting
1926
1927 #define SET_UP_TEST_FUNC                                        \
1928   using namespace AtomicBatchDispatcherTesting;                 \
1929   folly::EventBase evb;                                         \
1930   auto& executor = getFiberManager(evb);                        \
1931   const size_t COUNT = 11;                                      \
1932   std::vector<Job> jobs;                                        \
1933   jobs.reserve(COUNT);                                          \
1934   std::vector<folly::Optional<folly::Future<ResultT>>> results; \
1935   results.reserve(COUNT);                                       \
1936   DispatchFunctionT dispatchFunc
1937
1938 TEST(FiberManager, ABD_Test) {
1939   SET_UP_TEST_FUNC;
1940
1941   //
1942   // Testing AtomicBatchDispatcher with explicit call to commit()
1943   //
1944   dispatchFunc = userDispatchFunc;
1945   auto atomicBatchDispatcher =
1946       createAtomicBatchDispatcher(std::move(dispatchFunc));
1947   createJobs(atomicBatchDispatcher, jobs, COUNT);
1948   dispatchJobs(executor, jobs, results);
1949   atomicBatchDispatcher.commit();
1950   evb.loop();
1951   validateResults(results, COUNT);
1952 }
1953
1954 TEST(FiberManager, ABD_DispatcherDestroyedBeforeCallingCommit) {
1955   SET_UP_TEST_FUNC;
1956
1957   //
1958   // Testing AtomicBatchDispatcher destroyed before calling commit.
1959   // Handles error cases for:
1960   // - User might have forgotten to add the call to commit() in the code
1961   // - An unexpected exception got thrown in user code before commit() is called
1962   //
1963   try {
1964     dispatchFunc = userDispatchFunc;
1965     auto atomicBatchDispatcher =
1966         createAtomicBatchDispatcher(std::move(dispatchFunc));
1967     createJobs(atomicBatchDispatcher, jobs, COUNT);
1968     dispatchJobs(executor, jobs, results);
1969     throw std::runtime_error(
1970         "Unexpected exception in user code before commit called");
1971     // atomicBatchDispatcher.commit();
1972   } catch (...) {
1973     /* User code handles the exception and does not exit process */
1974   }
1975   evb.loop();
1976   validateResults<ABDCommitNotCalledException>(results, COUNT);
1977 }
1978
1979 TEST(FiberManager, ABD_PreprocessingFailureTest) {
1980   SET_UP_TEST_FUNC;
1981
1982   //
1983   // Testing preprocessing failure on a job throws
1984   //
1985   dispatchFunc = userDispatchFunc;
1986   auto atomicBatchDispatcher =
1987       createAtomicBatchDispatcher(std::move(dispatchFunc));
1988   createJobs(atomicBatchDispatcher, jobs, COUNT);
1989   dispatchJobs(executor, jobs, results, DispatchProblem::PreprocessThrows, 8);
1990   atomicBatchDispatcher.commit();
1991   evb.loop();
1992   validateResults<ABDTokenNotDispatchedException>(results, COUNT - 1);
1993 }
1994
1995 TEST(FiberManager, ABD_MultipleDispatchOnSameTokenErrorTest) {
1996   SET_UP_TEST_FUNC;
1997
1998   //
1999   // Testing that calling dispatch more than once on the same token throws
2000   //
2001   dispatchFunc = userDispatchFunc;
2002   auto atomicBatchDispatcher =
2003       createAtomicBatchDispatcher(std::move(dispatchFunc));
2004   createJobs(atomicBatchDispatcher, jobs, COUNT);
2005   dispatchJobs(executor, jobs, results, DispatchProblem::DuplicateDispatch, 4);
2006   atomicBatchDispatcher.commit();
2007   evb.loop();
2008 }
2009
2010 TEST(FiberManager, ABD_GetTokenCalledAfterCommitTest) {
2011   SET_UP_TEST_FUNC;
2012
2013   //
2014   // Testing that exception set on attempt to call getToken after commit called
2015   //
2016   dispatchFunc = userDispatchFunc;
2017   auto atomicBatchDispatcher =
2018       createAtomicBatchDispatcher(std::move(dispatchFunc));
2019   createJobs(atomicBatchDispatcher, jobs, COUNT);
2020   atomicBatchDispatcher.commit();
2021   EXPECT_THROW(atomicBatchDispatcher.getToken(), ABDUsageException);
2022   dispatchJobs(executor, jobs, results);
2023   EXPECT_THROW(atomicBatchDispatcher.getToken(), ABDUsageException);
2024   evb.loop();
2025   validateResults(results, COUNT);
2026   EXPECT_THROW(atomicBatchDispatcher.getToken(), ABDUsageException);
2027 }
2028
2029 TEST(FiberManager, ABD_UserProvidedBatchDispatchThrowsTest) {
2030   SET_UP_TEST_FUNC;
2031
2032   //
2033   // Testing that exception is set if user provided batch dispatch throws
2034   //
2035   dispatchFunc = [](std::vector<ValueT>&& inputs) -> std::vector<ResultT> {
2036     (void)userDispatchFunc(std::move(inputs));
2037     throw std::runtime_error("Unexpected exception in user dispatch function");
2038   };
2039   auto atomicBatchDispatcher =
2040       createAtomicBatchDispatcher(std::move(dispatchFunc));
2041   createJobs(atomicBatchDispatcher, jobs, COUNT);
2042   dispatchJobs(executor, jobs, results);
2043   atomicBatchDispatcher.commit();
2044   evb.loop();
2045   validateResults<std::runtime_error>(results, COUNT);
2046 }
2047
2048 TEST(FiberManager, VirtualEventBase) {
2049   bool done1{false};
2050   bool done2{false};
2051   {
2052     folly::ScopedEventBaseThread thread;
2053
2054     auto evb1 =
2055         std::make_unique<folly::VirtualEventBase>(*thread.getEventBase());
2056     auto& evb2 = thread.getEventBase()->getVirtualEventBase();
2057
2058     getFiberManager(*evb1).addTaskRemote([&] {
2059       Baton baton;
2060       baton.try_wait_for(std::chrono::milliseconds{100});
2061
2062       done1 = true;
2063     });
2064
2065     getFiberManager(evb2).addTaskRemote([&] {
2066       Baton baton;
2067       baton.try_wait_for(std::chrono::milliseconds{200});
2068
2069       done2 = true;
2070     });
2071
2072     EXPECT_FALSE(done1);
2073     EXPECT_FALSE(done2);
2074
2075     evb1.reset();
2076     EXPECT_TRUE(done1);
2077     EXPECT_FALSE(done2);
2078   }
2079   EXPECT_TRUE(done2);
2080 }
2081
2082 TEST(TimedMutex, ThreadsAndFibersDontDeadlock) {
2083   folly::EventBase evb;
2084   auto& fm = getFiberManager(evb);
2085   TimedMutex mutex;
2086   std::thread testThread([&] {
2087     for (int i = 0; i < 100; i++) {
2088       mutex.lock();
2089       mutex.unlock();
2090       {
2091         Baton b;
2092         b.try_wait_for(std::chrono::milliseconds(1));
2093       }
2094     }
2095   });
2096
2097   for (int numFibers = 0; numFibers < 100; numFibers++) {
2098     fm.addTask([&] {
2099       for (int i = 0; i < 20; i++) {
2100         mutex.lock();
2101         {
2102           Baton b;
2103           b.try_wait_for(std::chrono::milliseconds(1));
2104         }
2105         mutex.unlock();
2106         {
2107           Baton b;
2108           b.try_wait_for(std::chrono::milliseconds(1));
2109         }
2110       }
2111     });
2112   }
2113
2114   evb.loop();
2115   EXPECT_EQ(0, fm.hasTasks());
2116   testThread.join();
2117 }
2118
2119 TEST(TimedMutex, ThreadFiberDeadlockOrder) {
2120   folly::EventBase evb;
2121   auto& fm = getFiberManager(evb);
2122   TimedMutex mutex;
2123
2124   mutex.lock();
2125   std::thread unlockThread([&] {
2126     /* sleep override */ std::this_thread::sleep_for(
2127         std::chrono::milliseconds{100});
2128     mutex.unlock();
2129   });
2130
2131   fm.addTask([&] { std::lock_guard<TimedMutex> lg(mutex); });
2132   fm.addTask([&] {
2133     runInMainContext([&] {
2134       auto locked = mutex.timed_lock(std::chrono::seconds{1});
2135       EXPECT_TRUE(locked);
2136       if (locked) {
2137         mutex.unlock();
2138       }
2139     });
2140   });
2141
2142   evb.loopOnce();
2143   EXPECT_EQ(0, fm.hasTasks());
2144
2145   unlockThread.join();
2146 }
2147
2148 TEST(TimedMutex, ThreadFiberDeadlockRace) {
2149   folly::EventBase evb;
2150   auto& fm = getFiberManager(evb);
2151   TimedMutex mutex;
2152
2153   mutex.lock();
2154
2155   fm.addTask([&] {
2156     auto locked = mutex.timed_lock(std::chrono::seconds{1});
2157     EXPECT_TRUE(locked);
2158     if (locked) {
2159       mutex.unlock();
2160     }
2161   });
2162   fm.addTask([&] {
2163     mutex.unlock();
2164     runInMainContext([&] {
2165       auto locked = mutex.timed_lock(std::chrono::seconds{1});
2166       EXPECT_TRUE(locked);
2167       if (locked) {
2168         mutex.unlock();
2169       }
2170     });
2171   });
2172
2173   evb.loopOnce();
2174   EXPECT_EQ(0, fm.hasTasks());
2175 }
2176
2177 /**
2178  * Test that we can properly track fiber stack usage.
2179  *
2180  * This functionality can only be enabled when ASAN is disabled, so avoid
2181  * running this test with ASAN.
2182  */
2183 #ifndef FOLLY_SANITIZE_ADDRESS
2184 TEST(FiberManager, recordStack) {
2185   auto f = [] {
2186     folly::fibers::FiberManager::Options opts;
2187     opts.recordStackEvery = 1;
2188
2189     FiberManager fm(std::make_unique<SimpleLoopController>(), opts);
2190     auto& loopController =
2191         dynamic_cast<SimpleLoopController&>(fm.loopController());
2192
2193     static constexpr size_t n = 1000;
2194     int s = 0;
2195     fm.addTask([&]() {
2196       int b[n] = {0};
2197       for (size_t i = 0; i < n; ++i) {
2198         b[i] = i;
2199       }
2200       for (size_t i = 0; i + 1 < n; ++i) {
2201         s += b[i] * b[i + 1];
2202       }
2203     });
2204
2205     (void)s;
2206
2207     loopController.loop([&]() { loopController.stop(); });
2208
2209     // Check that we properly accounted fiber stack usage.
2210     EXPECT_LT(n * sizeof(int), fm.stackHighWatermark());
2211   };
2212   std::thread(f).join();
2213 }
2214 #endif