folly/fibers/test/FibersTest.cpp

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