folly/experimental/fibers/test/FibersTest.cpp

   1 /*
   2  * Copyright 2016 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 <gtest/gtest.h>
  21
  22 #include <folly/Benchmark.h>
  23 #include <folly/Memory.h>
  24 #include <folly/futures/Future.h>
  25
  26 #include <folly/experimental/fibers/AddTasks.h>
  27 #include <folly/experimental/fibers/EventBaseLoopController.h>
  28 #include <folly/experimental/fibers/FiberManager.h>
  29 #include <folly/experimental/fibers/FiberManagerMap.h>
  30 #include <folly/experimental/fibers/GenericBaton.h>
  31 #include <folly/experimental/fibers/SimpleLoopController.h>
  32 #include <folly/experimental/fibers/WhenN.h>
  33
  34 using namespace folly::fibers;
  35
  36 using folly::Try;
  37
  38 TEST(FiberManager, batonTimedWaitTimeout) {
  39   bool taskAdded = false;
  40   size_t iterations = 0;
  41
  42   FiberManager manager(folly::make_unique<SimpleLoopController>());
  43   auto& loopController =
  44       dynamic_cast<SimpleLoopController&>(manager.loopController());
  45
  46   auto loopFunc = [&]() {
  47     if (!taskAdded) {
  48       manager.addTask([&]() {
  49         Baton baton;
  50
  51         auto res = baton.timed_wait(std::chrono::milliseconds(230));
  52
  53         EXPECT_FALSE(res);
  54         EXPECT_EQ(5, iterations);
  55
  56         loopController.stop();
  57       });
  58       manager.addTask([&]() {
  59         Baton baton;
  60
  61         auto res = baton.timed_wait(std::chrono::milliseconds(130));
  62
  63         EXPECT_FALSE(res);
  64         EXPECT_EQ(3, iterations);
  65
  66         loopController.stop();
  67       });
  68       taskAdded = true;
  69     } else {
  70       std::this_thread::sleep_for(std::chrono::milliseconds(50));
  71       iterations++;
  72     }
  73   };
  74
  75   loopController.loop(std::move(loopFunc));
  76 }
  77
  78 TEST(FiberManager, batonTimedWaitPost) {
  79   bool taskAdded = false;
  80   size_t iterations = 0;
  81   Baton* baton_ptr;
  82
  83   FiberManager manager(folly::make_unique<SimpleLoopController>());
  84   auto& loopController =
  85       dynamic_cast<SimpleLoopController&>(manager.loopController());
  86
  87   auto loopFunc = [&]() {
  88     if (!taskAdded) {
  89       manager.addTask([&]() {
  90         Baton baton;
  91         baton_ptr = &baton;
  92
  93         auto res = baton.timed_wait(std::chrono::milliseconds(130));
  94
  95         EXPECT_TRUE(res);
  96         EXPECT_EQ(2, iterations);
  97
  98         loopController.stop();
  99       });
 100       taskAdded = true;
 101     } else {
 102       std::this_thread::sleep_for(std::chrono::milliseconds(50));
 103       iterations++;
 104       if (iterations == 2) {
 105         baton_ptr->post();
 106       }
 107     }
 108   };
 109
 110   loopController.loop(std::move(loopFunc));
 111 }
 112
 113 TEST(FiberManager, batonTimedWaitTimeoutEvb) {
 114   size_t tasksComplete = 0;
 115
 116   folly::EventBase evb;
 117
 118   FiberManager manager(folly::make_unique<EventBaseLoopController>());
 119   dynamic_cast<EventBaseLoopController&>(manager.loopController())
 120       .attachEventBase(evb);
 121
 122   auto task = [&](size_t timeout_ms) {
 123     Baton baton;
 124
 125     auto start = EventBaseLoopController::Clock::now();
 126     auto res = baton.timed_wait(std::chrono::milliseconds(timeout_ms));
 127     auto finish = EventBaseLoopController::Clock::now();
 128
 129     EXPECT_FALSE(res);
 130
 131     auto duration_ms =
 132         std::chrono::duration_cast<std::chrono::milliseconds>(finish - start);
 133
 134     EXPECT_GT(duration_ms.count(), timeout_ms - 50);
 135     EXPECT_LT(duration_ms.count(), timeout_ms + 50);
 136
 137     if (++tasksComplete == 2) {
 138       evb.terminateLoopSoon();
 139     }
 140   };
 141
 142   evb.runInEventBaseThread([&]() {
 143     manager.addTask([&]() { task(500); });
 144     manager.addTask([&]() { task(250); });
 145   });
 146
 147   evb.loopForever();
 148
 149   EXPECT_EQ(2, tasksComplete);
 150 }
 151
 152 TEST(FiberManager, batonTimedWaitPostEvb) {
 153   size_t tasksComplete = 0;
 154
 155   folly::EventBase evb;
 156
 157   FiberManager manager(folly::make_unique<EventBaseLoopController>());
 158   dynamic_cast<EventBaseLoopController&>(manager.loopController())
 159       .attachEventBase(evb);
 160
 161   evb.runInEventBaseThread([&]() {
 162     manager.addTask([&]() {
 163       Baton baton;
 164
 165       evb.tryRunAfterDelay([&]() { baton.post(); }, 100);
 166
 167       auto start = EventBaseLoopController::Clock::now();
 168       auto res = baton.timed_wait(std::chrono::milliseconds(130));
 169       auto finish = EventBaseLoopController::Clock::now();
 170
 171       EXPECT_TRUE(res);
 172
 173       auto duration_ms =
 174           std::chrono::duration_cast<std::chrono::milliseconds>(finish - start);
 175
 176       EXPECT_TRUE(duration_ms.count() > 95 && duration_ms.count() < 110);
 177
 178       if (++tasksComplete == 1) {
 179         evb.terminateLoopSoon();
 180       }
 181     });
 182   });
 183
 184   evb.loopForever();
 185
 186   EXPECT_EQ(1, tasksComplete);
 187 }
 188
 189 TEST(FiberManager, batonTryWait) {
 190   FiberManager manager(folly::make_unique<SimpleLoopController>());
 191
 192   // Check if try_wait and post work as expected
 193   Baton b;
 194
 195   manager.addTask([&]() {
 196     while (!b.try_wait()) {
 197     }
 198   });
 199   auto thr = std::thread([&]() {
 200     std::this_thread::sleep_for(std::chrono::milliseconds(300));
 201     b.post();
 202   });
 203
 204   manager.loopUntilNoReady();
 205   thr.join();
 206
 207   Baton c;
 208
 209   // Check try_wait without post
 210   manager.addTask([&]() {
 211     int cnt = 100;
 212     while (cnt && !c.try_wait()) {
 213       cnt--;
 214     }
 215     EXPECT_TRUE(!c.try_wait()); // must still hold
 216     EXPECT_EQ(cnt, 0);
 217   });
 218
 219   manager.loopUntilNoReady();
 220 }
 221
 222 TEST(FiberManager, genericBatonFiberWait) {
 223   FiberManager manager(folly::make_unique<SimpleLoopController>());
 224
 225   GenericBaton b;
 226   bool fiberRunning = false;
 227
 228   manager.addTask([&]() {
 229     EXPECT_EQ(manager.hasActiveFiber(), true);
 230     fiberRunning = true;
 231     b.wait();
 232     fiberRunning = false;
 233   });
 234
 235   EXPECT_FALSE(fiberRunning);
 236   manager.loopUntilNoReady();
 237   EXPECT_TRUE(fiberRunning); // ensure fiber still active
 238
 239   auto thr = std::thread([&]() {
 240     std::this_thread::sleep_for(std::chrono::milliseconds(300));
 241     b.post();
 242   });
 243
 244   while (fiberRunning) {
 245     manager.loopUntilNoReady();
 246   }
 247
 248   thr.join();
 249 }
 250
 251 TEST(FiberManager, genericBatonThreadWait) {
 252   FiberManager manager(folly::make_unique<SimpleLoopController>());
 253   GenericBaton b;
 254   std::atomic<bool> threadWaiting(false);
 255
 256   auto thr = std::thread([&]() {
 257     threadWaiting = true;
 258     b.wait();
 259     threadWaiting = false;
 260   });
 261
 262   while (!threadWaiting) {
 263   }
 264   std::this_thread::sleep_for(std::chrono::milliseconds(300));
 265
 266   manager.addTask([&]() {
 267     EXPECT_EQ(manager.hasActiveFiber(), true);
 268     EXPECT_TRUE(threadWaiting);
 269     b.post();
 270     while (threadWaiting) {
 271     }
 272   });
 273
 274   manager.loopUntilNoReady();
 275   thr.join();
 276 }
 277
 278 TEST(FiberManager, addTasksNoncopyable) {
 279   std::vector<Promise<int>> pendingFibers;
 280   bool taskAdded = false;
 281
 282   FiberManager manager(folly::make_unique<SimpleLoopController>());
 283   auto& loopController =
 284       dynamic_cast<SimpleLoopController&>(manager.loopController());
 285
 286   auto loopFunc = [&]() {
 287     if (!taskAdded) {
 288       manager.addTask([&]() {
 289         std::vector<std::function<std::unique_ptr<int>()>> funcs;
 290         for (size_t i = 0; i < 3; ++i) {
 291           funcs.push_back([i, &pendingFibers]() {
 292             await([&pendingFibers](Promise<int> promise) {
 293               pendingFibers.push_back(std::move(promise));
 294             });
 295             return folly::make_unique<int>(i * 2 + 1);
 296           });
 297         }
 298
 299         auto iter = addTasks(funcs.begin(), funcs.end());
 300
 301         size_t n = 0;
 302         while (iter.hasNext()) {
 303           auto result = iter.awaitNext();
 304           EXPECT_EQ(2 * iter.getTaskID() + 1, *result);
 305           EXPECT_GE(2 - n, pendingFibers.size());
 306           ++n;
 307         }
 308         EXPECT_EQ(3, n);
 309       });
 310       taskAdded = true;
 311     } else if (pendingFibers.size()) {
 312       pendingFibers.back().setValue(0);
 313       pendingFibers.pop_back();
 314     } else {
 315       loopController.stop();
 316     }
 317   };
 318
 319   loopController.loop(std::move(loopFunc));
 320 }
 321
 322 TEST(FiberManager, awaitThrow) {
 323   folly::EventBase evb;
 324   struct ExpectedException {};
 325   getFiberManager(evb)
 326       .addTaskFuture([&] {
 327         EXPECT_THROW(
 328           await([](Promise<int> p) {
 329               p.setValue(42);
 330               throw ExpectedException();
 331             }),
 332           ExpectedException
 333         );
 334
 335         EXPECT_THROW(
 336           await([&](Promise<int> p) {
 337               evb.runInEventBaseThread([p = std::move(p)]() mutable {
 338                   p.setValue(42);
 339                 });
 340               throw ExpectedException();
 341             }),
 342           ExpectedException);
 343       })
 344       .waitVia(&evb);
 345 }
 346
 347 TEST(FiberManager, addTasksThrow) {
 348   std::vector<Promise<int>> pendingFibers;
 349   bool taskAdded = false;
 350
 351   FiberManager manager(folly::make_unique<SimpleLoopController>());
 352   auto& loopController =
 353       dynamic_cast<SimpleLoopController&>(manager.loopController());
 354
 355   auto loopFunc = [&]() {
 356     if (!taskAdded) {
 357       manager.addTask([&]() {
 358         std::vector<std::function<int()>> funcs;
 359         for (size_t i = 0; i < 3; ++i) {
 360           funcs.push_back([i, &pendingFibers]() {
 361             await([&pendingFibers](Promise<int> promise) {
 362               pendingFibers.push_back(std::move(promise));
 363             });
 364             if (i % 2 == 0) {
 365               throw std::runtime_error("Runtime");
 366             }
 367             return i * 2 + 1;
 368           });
 369         }
 370
 371         auto iter = addTasks(funcs.begin(), funcs.end());
 372
 373         size_t n = 0;
 374         while (iter.hasNext()) {
 375           try {
 376             int result = iter.awaitNext();
 377             EXPECT_EQ(1, iter.getTaskID() % 2);
 378             EXPECT_EQ(2 * iter.getTaskID() + 1, result);
 379           } catch (...) {
 380             EXPECT_EQ(0, iter.getTaskID() % 2);
 381           }
 382           EXPECT_GE(2 - n, pendingFibers.size());
 383           ++n;
 384         }
 385         EXPECT_EQ(3, n);
 386       });
 387       taskAdded = true;
 388     } else if (pendingFibers.size()) {
 389       pendingFibers.back().setValue(0);
 390       pendingFibers.pop_back();
 391     } else {
 392       loopController.stop();
 393     }
 394   };
 395
 396   loopController.loop(std::move(loopFunc));
 397 }
 398
 399 TEST(FiberManager, addTasksVoid) {
 400   std::vector<Promise<int>> pendingFibers;
 401   bool taskAdded = false;
 402
 403   FiberManager manager(folly::make_unique<SimpleLoopController>());
 404   auto& loopController =
 405       dynamic_cast<SimpleLoopController&>(manager.loopController());
 406
 407   auto loopFunc = [&]() {
 408     if (!taskAdded) {
 409       manager.addTask([&]() {
 410         std::vector<std::function<void()>> funcs;
 411         for (size_t i = 0; i < 3; ++i) {
 412           funcs.push_back([i, &pendingFibers]() {
 413             await([&pendingFibers](Promise<int> promise) {
 414               pendingFibers.push_back(std::move(promise));
 415             });
 416           });
 417         }
 418
 419         auto iter = addTasks(funcs.begin(), funcs.end());
 420
 421         size_t n = 0;
 422         while (iter.hasNext()) {
 423           iter.awaitNext();
 424           EXPECT_GE(2 - n, pendingFibers.size());
 425           ++n;
 426         }
 427         EXPECT_EQ(3, n);
 428       });
 429       taskAdded = true;
 430     } else if (pendingFibers.size()) {
 431       pendingFibers.back().setValue(0);
 432       pendingFibers.pop_back();
 433     } else {
 434       loopController.stop();
 435     }
 436   };
 437
 438   loopController.loop(std::move(loopFunc));
 439 }
 440
 441 TEST(FiberManager, addTasksVoidThrow) {
 442   std::vector<Promise<int>> pendingFibers;
 443   bool taskAdded = false;
 444
 445   FiberManager manager(folly::make_unique<SimpleLoopController>());
 446   auto& loopController =
 447       dynamic_cast<SimpleLoopController&>(manager.loopController());
 448
 449   auto loopFunc = [&]() {
 450     if (!taskAdded) {
 451       manager.addTask([&]() {
 452         std::vector<std::function<void()>> funcs;
 453         for (size_t i = 0; i < 3; ++i) {
 454           funcs.push_back([i, &pendingFibers]() {
 455             await([&pendingFibers](Promise<int> promise) {
 456               pendingFibers.push_back(std::move(promise));
 457             });
 458             if (i % 2 == 0) {
 459               throw std::runtime_error("");
 460             }
 461           });
 462         }
 463
 464         auto iter = addTasks(funcs.begin(), funcs.end());
 465
 466         size_t n = 0;
 467         while (iter.hasNext()) {
 468           try {
 469             iter.awaitNext();
 470             EXPECT_EQ(1, iter.getTaskID() % 2);
 471           } catch (...) {
 472             EXPECT_EQ(0, iter.getTaskID() % 2);
 473           }
 474           EXPECT_GE(2 - n, pendingFibers.size());
 475           ++n;
 476         }
 477         EXPECT_EQ(3, n);
 478       });
 479       taskAdded = true;
 480     } else if (pendingFibers.size()) {
 481       pendingFibers.back().setValue(0);
 482       pendingFibers.pop_back();
 483     } else {
 484       loopController.stop();
 485     }
 486   };
 487
 488   loopController.loop(std::move(loopFunc));
 489 }
 490
 491 TEST(FiberManager, addTasksReserve) {
 492   std::vector<Promise<int>> pendingFibers;
 493   bool taskAdded = false;
 494
 495   FiberManager manager(folly::make_unique<SimpleLoopController>());
 496   auto& loopController =
 497       dynamic_cast<SimpleLoopController&>(manager.loopController());
 498
 499   auto loopFunc = [&]() {
 500     if (!taskAdded) {
 501       manager.addTask([&]() {
 502         std::vector<std::function<void()>> funcs;
 503         for (size_t i = 0; i < 3; ++i) {
 504           funcs.push_back([&pendingFibers]() {
 505             await([&pendingFibers](Promise<int> promise) {
 506               pendingFibers.push_back(std::move(promise));
 507             });
 508           });
 509         }
 510
 511         auto iter = addTasks(funcs.begin(), funcs.end());
 512
 513         iter.reserve(2);
 514         EXPECT_TRUE(iter.hasCompleted());
 515         EXPECT_TRUE(iter.hasPending());
 516         EXPECT_TRUE(iter.hasNext());
 517
 518         iter.awaitNext();
 519         EXPECT_TRUE(iter.hasCompleted());
 520         EXPECT_TRUE(iter.hasPending());
 521         EXPECT_TRUE(iter.hasNext());
 522
 523         iter.awaitNext();
 524         EXPECT_FALSE(iter.hasCompleted());
 525         EXPECT_TRUE(iter.hasPending());
 526         EXPECT_TRUE(iter.hasNext());
 527
 528         iter.awaitNext();
 529         EXPECT_FALSE(iter.hasCompleted());
 530         EXPECT_FALSE(iter.hasPending());
 531         EXPECT_FALSE(iter.hasNext());
 532       });
 533       taskAdded = true;
 534     } else if (pendingFibers.size()) {
 535       pendingFibers.back().setValue(0);
 536       pendingFibers.pop_back();
 537     } else {
 538       loopController.stop();
 539     }
 540   };
 541
 542   loopController.loop(std::move(loopFunc));
 543 }
 544
 545 TEST(FiberManager, addTaskDynamic) {
 546   folly::EventBase evb;
 547
 548   Baton batons[3];
 549
 550   auto makeTask = [&](size_t taskId) {
 551     return [&, taskId]() -> size_t {
 552       batons[taskId].wait();
 553       return taskId;
 554     };
 555   };
 556
 557   getFiberManager(evb)
 558       .addTaskFuture([&]() {
 559         TaskIterator<size_t> iterator;
 560
 561         iterator.addTask(makeTask(0));
 562         iterator.addTask(makeTask(1));
 563
 564         batons[1].post();
 565
 566         EXPECT_EQ(1, iterator.awaitNext());
 567
 568         iterator.addTask(makeTask(2));
 569
 570         batons[2].post();
 571
 572         EXPECT_EQ(2, iterator.awaitNext());
 573
 574         batons[0].post();
 575
 576         EXPECT_EQ(0, iterator.awaitNext());
 577       })
 578       .waitVia(&evb);
 579 }
 580
 581 TEST(FiberManager, forEach) {
 582   std::vector<Promise<int>> pendingFibers;
 583   bool taskAdded = false;
 584
 585   FiberManager manager(folly::make_unique<SimpleLoopController>());
 586   auto& loopController =
 587       dynamic_cast<SimpleLoopController&>(manager.loopController());
 588
 589   auto loopFunc = [&]() {
 590     if (!taskAdded) {
 591       manager.addTask([&]() {
 592         std::vector<std::function<int()>> funcs;
 593         for (size_t i = 0; i < 3; ++i) {
 594           funcs.push_back([i, &pendingFibers]() {
 595             await([&pendingFibers](Promise<int> promise) {
 596               pendingFibers.push_back(std::move(promise));
 597             });
 598             return i * 2 + 1;
 599           });
 600         }
 601
 602         std::vector<std::pair<size_t, int>> results;
 603         forEach(funcs.begin(), funcs.end(), [&results](size_t id, int result) {
 604           results.emplace_back(id, result);
 605         });
 606         EXPECT_EQ(3, results.size());
 607         EXPECT_TRUE(pendingFibers.empty());
 608         for (size_t i = 0; i < 3; ++i) {
 609           EXPECT_EQ(results[i].first * 2 + 1, results[i].second);
 610         }
 611       });
 612       taskAdded = true;
 613     } else if (pendingFibers.size()) {
 614       pendingFibers.back().setValue(0);
 615       pendingFibers.pop_back();
 616     } else {
 617       loopController.stop();
 618     }
 619   };
 620
 621   loopController.loop(std::move(loopFunc));
 622 }
 623
 624 TEST(FiberManager, collectN) {
 625   std::vector<Promise<int>> pendingFibers;
 626   bool taskAdded = false;
 627
 628   FiberManager manager(folly::make_unique<SimpleLoopController>());
 629   auto& loopController =
 630       dynamic_cast<SimpleLoopController&>(manager.loopController());
 631
 632   auto loopFunc = [&]() {
 633     if (!taskAdded) {
 634       manager.addTask([&]() {
 635         std::vector<std::function<int()>> funcs;
 636         for (size_t i = 0; i < 3; ++i) {
 637           funcs.push_back([i, &pendingFibers]() {
 638             await([&pendingFibers](Promise<int> promise) {
 639               pendingFibers.push_back(std::move(promise));
 640             });
 641             return i * 2 + 1;
 642           });
 643         }
 644
 645         auto results = collectN(funcs.begin(), funcs.end(), 2);
 646         EXPECT_EQ(2, results.size());
 647         EXPECT_EQ(1, pendingFibers.size());
 648         for (size_t i = 0; i < 2; ++i) {
 649           EXPECT_EQ(results[i].first * 2 + 1, results[i].second);
 650         }
 651       });
 652       taskAdded = true;
 653     } else if (pendingFibers.size()) {
 654       pendingFibers.back().setValue(0);
 655       pendingFibers.pop_back();
 656     } else {
 657       loopController.stop();
 658     }
 659   };
 660
 661   loopController.loop(std::move(loopFunc));
 662 }
 663
 664 TEST(FiberManager, collectNThrow) {
 665   std::vector<Promise<int>> pendingFibers;
 666   bool taskAdded = false;
 667
 668   FiberManager manager(folly::make_unique<SimpleLoopController>());
 669   auto& loopController =
 670       dynamic_cast<SimpleLoopController&>(manager.loopController());
 671
 672   auto loopFunc = [&]() {
 673     if (!taskAdded) {
 674       manager.addTask([&]() {
 675         std::vector<std::function<int()>> funcs;
 676         for (size_t i = 0; i < 3; ++i) {
 677           funcs.push_back([i, &pendingFibers]() {
 678             await([&pendingFibers](Promise<int> promise) {
 679               pendingFibers.push_back(std::move(promise));
 680             });
 681             throw std::runtime_error("Runtime");
 682             return i * 2 + 1;
 683           });
 684         }
 685
 686         try {
 687           collectN(funcs.begin(), funcs.end(), 2);
 688         } catch (...) {
 689           EXPECT_EQ(1, pendingFibers.size());
 690         }
 691       });
 692       taskAdded = true;
 693     } else if (pendingFibers.size()) {
 694       pendingFibers.back().setValue(0);
 695       pendingFibers.pop_back();
 696     } else {
 697       loopController.stop();
 698     }
 699   };
 700
 701   loopController.loop(std::move(loopFunc));
 702 }
 703
 704 TEST(FiberManager, collectNVoid) {
 705   std::vector<Promise<int>> pendingFibers;
 706   bool taskAdded = false;
 707
 708   FiberManager manager(folly::make_unique<SimpleLoopController>());
 709   auto& loopController =
 710       dynamic_cast<SimpleLoopController&>(manager.loopController());
 711
 712   auto loopFunc = [&]() {
 713     if (!taskAdded) {
 714       manager.addTask([&]() {
 715         std::vector<std::function<void()>> funcs;
 716         for (size_t i = 0; i < 3; ++i) {
 717           funcs.push_back([i, &pendingFibers]() {
 718             await([&pendingFibers](Promise<int> promise) {
 719               pendingFibers.push_back(std::move(promise));
 720             });
 721           });
 722         }
 723
 724         auto results = collectN(funcs.begin(), funcs.end(), 2);
 725         EXPECT_EQ(2, results.size());
 726         EXPECT_EQ(1, pendingFibers.size());
 727       });
 728       taskAdded = true;
 729     } else if (pendingFibers.size()) {
 730       pendingFibers.back().setValue(0);
 731       pendingFibers.pop_back();
 732     } else {
 733       loopController.stop();
 734     }
 735   };
 736
 737   loopController.loop(std::move(loopFunc));
 738 }
 739
 740 TEST(FiberManager, collectNVoidThrow) {
 741   std::vector<Promise<int>> pendingFibers;
 742   bool taskAdded = false;
 743
 744   FiberManager manager(folly::make_unique<SimpleLoopController>());
 745   auto& loopController =
 746       dynamic_cast<SimpleLoopController&>(manager.loopController());
 747
 748   auto loopFunc = [&]() {
 749     if (!taskAdded) {
 750       manager.addTask([&]() {
 751         std::vector<std::function<void()>> funcs;
 752         for (size_t i = 0; i < 3; ++i) {
 753           funcs.push_back([i, &pendingFibers]() {
 754             await([&pendingFibers](Promise<int> promise) {
 755               pendingFibers.push_back(std::move(promise));
 756             });
 757             throw std::runtime_error("Runtime");
 758           });
 759         }
 760
 761         try {
 762           collectN(funcs.begin(), funcs.end(), 2);
 763         } catch (...) {
 764           EXPECT_EQ(1, pendingFibers.size());
 765         }
 766       });
 767       taskAdded = true;
 768     } else if (pendingFibers.size()) {
 769       pendingFibers.back().setValue(0);
 770       pendingFibers.pop_back();
 771     } else {
 772       loopController.stop();
 773     }
 774   };
 775
 776   loopController.loop(std::move(loopFunc));
 777 }
 778
 779 TEST(FiberManager, collectAll) {
 780   std::vector<Promise<int>> pendingFibers;
 781   bool taskAdded = false;
 782
 783   FiberManager manager(folly::make_unique<SimpleLoopController>());
 784   auto& loopController =
 785       dynamic_cast<SimpleLoopController&>(manager.loopController());
 786
 787   auto loopFunc = [&]() {
 788     if (!taskAdded) {
 789       manager.addTask([&]() {
 790         std::vector<std::function<int()>> funcs;
 791         for (size_t i = 0; i < 3; ++i) {
 792           funcs.push_back([i, &pendingFibers]() {
 793             await([&pendingFibers](Promise<int> promise) {
 794               pendingFibers.push_back(std::move(promise));
 795             });
 796             return i * 2 + 1;
 797           });
 798         }
 799
 800         auto results = collectAll(funcs.begin(), funcs.end());
 801         EXPECT_TRUE(pendingFibers.empty());
 802         for (size_t i = 0; i < 3; ++i) {
 803           EXPECT_EQ(i * 2 + 1, results[i]);
 804         }
 805       });
 806       taskAdded = true;
 807     } else if (pendingFibers.size()) {
 808       pendingFibers.back().setValue(0);
 809       pendingFibers.pop_back();
 810     } else {
 811       loopController.stop();
 812     }
 813   };
 814
 815   loopController.loop(std::move(loopFunc));
 816 }
 817
 818 TEST(FiberManager, collectAllVoid) {
 819   std::vector<Promise<int>> pendingFibers;
 820   bool taskAdded = false;
 821
 822   FiberManager manager(folly::make_unique<SimpleLoopController>());
 823   auto& loopController =
 824       dynamic_cast<SimpleLoopController&>(manager.loopController());
 825
 826   auto loopFunc = [&]() {
 827     if (!taskAdded) {
 828       manager.addTask([&]() {
 829         std::vector<std::function<void()>> funcs;
 830         for (size_t i = 0; i < 3; ++i) {
 831           funcs.push_back([i, &pendingFibers]() {
 832             await([&pendingFibers](Promise<int> promise) {
 833               pendingFibers.push_back(std::move(promise));
 834             });
 835           });
 836         }
 837
 838         collectAll(funcs.begin(), funcs.end());
 839         EXPECT_TRUE(pendingFibers.empty());
 840       });
 841       taskAdded = true;
 842     } else if (pendingFibers.size()) {
 843       pendingFibers.back().setValue(0);
 844       pendingFibers.pop_back();
 845     } else {
 846       loopController.stop();
 847     }
 848   };
 849
 850   loopController.loop(std::move(loopFunc));
 851 }
 852
 853 TEST(FiberManager, collectAny) {
 854   std::vector<Promise<int>> pendingFibers;
 855   bool taskAdded = false;
 856
 857   FiberManager manager(folly::make_unique<SimpleLoopController>());
 858   auto& loopController =
 859       dynamic_cast<SimpleLoopController&>(manager.loopController());
 860
 861   auto loopFunc = [&]() {
 862     if (!taskAdded) {
 863       manager.addTask([&]() {
 864         std::vector<std::function<int()>> funcs;
 865         for (size_t i = 0; i < 3; ++i) {
 866           funcs.push_back([i, &pendingFibers]() {
 867             await([&pendingFibers](Promise<int> promise) {
 868               pendingFibers.push_back(std::move(promise));
 869             });
 870             if (i == 1) {
 871               throw std::runtime_error("This exception will be ignored");
 872             }
 873             return i * 2 + 1;
 874           });
 875         }
 876
 877         auto result = collectAny(funcs.begin(), funcs.end());
 878         EXPECT_EQ(2, pendingFibers.size());
 879         EXPECT_EQ(2, result.first);
 880         EXPECT_EQ(2 * 2 + 1, result.second);
 881       });
 882       taskAdded = true;
 883     } else if (pendingFibers.size()) {
 884       pendingFibers.back().setValue(0);
 885       pendingFibers.pop_back();
 886     } else {
 887       loopController.stop();
 888     }
 889   };
 890
 891   loopController.loop(std::move(loopFunc));
 892 }
 893
 894 namespace {
 895 /* Checks that this function was run from a main context,
 896    by comparing an address on a stack to a known main stack address
 897    and a known related fiber stack address.  The assumption
 898    is that fiber stack and main stack will be far enough apart,
 899    while any two values on the same stack will be close. */
 900 void expectMainContext(bool& ran, int* mainLocation, int* fiberLocation) {
 901   int here;
 902   /* 2 pages is a good guess */
 903   constexpr ssize_t DISTANCE = 0x2000 / sizeof(int);
 904   if (fiberLocation) {
 905     EXPECT_TRUE(std::abs(&here - fiberLocation) > DISTANCE);
 906   }
 907   if (mainLocation) {
 908     EXPECT_TRUE(std::abs(&here - mainLocation) < DISTANCE);
 909   }
 910
 911   EXPECT_FALSE(ran);
 912   ran = true;
 913 }
 914 }
 915
 916 TEST(FiberManager, runInMainContext) {
 917   FiberManager manager(folly::make_unique<SimpleLoopController>());
 918   auto& loopController =
 919       dynamic_cast<SimpleLoopController&>(manager.loopController());
 920
 921   bool checkRan = false;
 922
 923   int mainLocation;
 924   manager.runInMainContext(
 925       [&]() { expectMainContext(checkRan, &mainLocation, nullptr); });
 926   EXPECT_TRUE(checkRan);
 927
 928   checkRan = false;
 929
 930   manager.addTask([&]() {
 931     struct A {
 932       explicit A(int value_) : value(value_) {}
 933       A(const A&) = delete;
 934       A(A&&) = default;
 935
 936       int value;
 937     };
 938     int stackLocation;
 939     auto ret = runInMainContext([&]() {
 940       expectMainContext(checkRan, &mainLocation, &stackLocation);
 941       return A(42);
 942     });
 943     EXPECT_TRUE(checkRan);
 944     EXPECT_EQ(42, ret.value);
 945   });
 946
 947   loopController.loop([&]() { loopController.stop(); });
 948
 949   EXPECT_TRUE(checkRan);
 950 }
 951
 952 TEST(FiberManager, addTaskFinally) {
 953   FiberManager manager(folly::make_unique<SimpleLoopController>());
 954   auto& loopController =
 955       dynamic_cast<SimpleLoopController&>(manager.loopController());
 956
 957   bool checkRan = false;
 958
 959   int mainLocation;
 960
 961   manager.addTaskFinally(
 962       [&]() { return 1234; },
 963       [&](Try<int>&& result) {
 964         EXPECT_EQ(result.value(), 1234);
 965
 966         expectMainContext(checkRan, &mainLocation, nullptr);
 967       });
 968
 969   EXPECT_FALSE(checkRan);
 970
 971   loopController.loop([&]() { loopController.stop(); });
 972
 973   EXPECT_TRUE(checkRan);
 974 }
 975
 976 TEST(FiberManager, fibersPoolWithinLimit) {
 977   FiberManager::Options opts;
 978   opts.maxFibersPoolSize = 5;
 979
 980   FiberManager manager(folly::make_unique<SimpleLoopController>(), opts);
 981   auto& loopController =
 982       dynamic_cast<SimpleLoopController&>(manager.loopController());
 983
 984   size_t fibersRun = 0;
 985
 986   for (size_t i = 0; i < 5; ++i) {
 987     manager.addTask([&]() { ++fibersRun; });
 988   }
 989   loopController.loop([&]() { loopController.stop(); });
 990
 991   EXPECT_EQ(5, fibersRun);
 992   EXPECT_EQ(5, manager.fibersAllocated());
 993   EXPECT_EQ(5, manager.fibersPoolSize());
 994
 995   for (size_t i = 0; i < 5; ++i) {
 996     manager.addTask([&]() { ++fibersRun; });
 997   }
 998   loopController.loop([&]() { loopController.stop(); });
 999
1000   EXPECT_EQ(10, fibersRun);
1001   EXPECT_EQ(5, manager.fibersAllocated());
1002   EXPECT_EQ(5, manager.fibersPoolSize());
1003 }
1004
1005 TEST(FiberManager, fibersPoolOverLimit) {
1006   FiberManager::Options opts;
1007   opts.maxFibersPoolSize = 5;
1008
1009   FiberManager manager(folly::make_unique<SimpleLoopController>(), opts);
1010   auto& loopController =
1011       dynamic_cast<SimpleLoopController&>(manager.loopController());
1012
1013   size_t fibersRun = 0;
1014
1015   for (size_t i = 0; i < 10; ++i) {
1016     manager.addTask([&]() { ++fibersRun; });
1017   }
1018
1019   EXPECT_EQ(0, fibersRun);
1020   EXPECT_EQ(10, manager.fibersAllocated());
1021   EXPECT_EQ(0, manager.fibersPoolSize());
1022
1023   loopController.loop([&]() { loopController.stop(); });
1024
1025   EXPECT_EQ(10, fibersRun);
1026   EXPECT_EQ(5, manager.fibersAllocated());
1027   EXPECT_EQ(5, manager.fibersPoolSize());
1028 }
1029
1030 TEST(FiberManager, remoteFiberBasic) {
1031   FiberManager manager(folly::make_unique<SimpleLoopController>());
1032   auto& loopController =
1033       dynamic_cast<SimpleLoopController&>(manager.loopController());
1034
1035   int result[2];
1036   result[0] = result[1] = 0;
1037   folly::Optional<Promise<int>> savedPromise[2];
1038   manager.addTask([&]() {
1039     result[0] = await(
1040         [&](Promise<int> promise) { savedPromise[0] = std::move(promise); });
1041   });
1042   manager.addTask([&]() {
1043     result[1] = await(
1044         [&](Promise<int> promise) { savedPromise[1] = std::move(promise); });
1045   });
1046
1047   manager.loopUntilNoReady();
1048
1049   EXPECT_TRUE(savedPromise[0].hasValue());
1050   EXPECT_TRUE(savedPromise[1].hasValue());
1051   EXPECT_EQ(0, result[0]);
1052   EXPECT_EQ(0, result[1]);
1053
1054   std::thread remoteThread0{[&]() { savedPromise[0]->setValue(42); }};
1055   std::thread remoteThread1{[&]() { savedPromise[1]->setValue(43); }};
1056   remoteThread0.join();
1057   remoteThread1.join();
1058   EXPECT_EQ(0, result[0]);
1059   EXPECT_EQ(0, result[1]);
1060   /* Should only have scheduled once */
1061   EXPECT_EQ(1, loopController.remoteScheduleCalled());
1062
1063   manager.loopUntilNoReady();
1064   EXPECT_EQ(42, result[0]);
1065   EXPECT_EQ(43, result[1]);
1066 }
1067
1068 TEST(FiberManager, addTaskRemoteBasic) {
1069   FiberManager manager(folly::make_unique<SimpleLoopController>());
1070
1071   int result[2];
1072   result[0] = result[1] = 0;
1073   folly::Optional<Promise<int>> savedPromise[2];
1074
1075   std::thread remoteThread0{[&]() {
1076     manager.addTaskRemote([&]() {
1077       result[0] = await(
1078           [&](Promise<int> promise) { savedPromise[0] = std::move(promise); });
1079     });
1080   }};
1081   std::thread remoteThread1{[&]() {
1082     manager.addTaskRemote([&]() {
1083       result[1] = await(
1084           [&](Promise<int> promise) { savedPromise[1] = std::move(promise); });
1085     });
1086   }};
1087   remoteThread0.join();
1088   remoteThread1.join();
1089
1090   manager.loopUntilNoReady();
1091
1092   EXPECT_TRUE(savedPromise[0].hasValue());
1093   EXPECT_TRUE(savedPromise[1].hasValue());
1094   EXPECT_EQ(0, result[0]);
1095   EXPECT_EQ(0, result[1]);
1096
1097   savedPromise[0]->setValue(42);
1098   savedPromise[1]->setValue(43);
1099
1100   EXPECT_EQ(0, result[0]);
1101   EXPECT_EQ(0, result[1]);
1102
1103   manager.loopUntilNoReady();
1104   EXPECT_EQ(42, result[0]);
1105   EXPECT_EQ(43, result[1]);
1106 }
1107
1108 TEST(FiberManager, remoteHasTasks) {
1109   size_t counter = 0;
1110   FiberManager fm(folly::make_unique<SimpleLoopController>());
1111   std::thread remote([&]() { fm.addTaskRemote([&]() { ++counter; }); });
1112
1113   remote.join();
1114
1115   while (fm.hasTasks()) {
1116     fm.loopUntilNoReady();
1117   }
1118
1119   EXPECT_FALSE(fm.hasTasks());
1120   EXPECT_EQ(counter, 1);
1121 }
1122
1123 TEST(FiberManager, remoteHasReadyTasks) {
1124   int result = 0;
1125   folly::Optional<Promise<int>> savedPromise;
1126   FiberManager fm(folly::make_unique<SimpleLoopController>());
1127   std::thread remote([&]() {
1128     fm.addTaskRemote([&]() {
1129       result = await(
1130           [&](Promise<int> promise) { savedPromise = std::move(promise); });
1131       EXPECT_TRUE(fm.hasTasks());
1132     });
1133   });
1134
1135   remote.join();
1136   EXPECT_TRUE(fm.hasTasks());
1137
1138   fm.loopUntilNoReady();
1139   EXPECT_TRUE(fm.hasTasks());
1140
1141   std::thread remote2([&]() { savedPromise->setValue(47); });
1142   remote2.join();
1143   EXPECT_TRUE(fm.hasTasks());
1144
1145   fm.loopUntilNoReady();
1146   EXPECT_FALSE(fm.hasTasks());
1147
1148   EXPECT_EQ(result, 47);
1149 }
1150
1151 template <typename Data>
1152 void testFiberLocal() {
1153   FiberManager fm(
1154       LocalType<Data>(), folly::make_unique<SimpleLoopController>());
1155
1156   fm.addTask([]() {
1157     EXPECT_EQ(42, local<Data>().value);
1158
1159     local<Data>().value = 43;
1160
1161     addTask([]() {
1162       EXPECT_EQ(43, local<Data>().value);
1163
1164       local<Data>().value = 44;
1165
1166       addTask([]() { EXPECT_EQ(44, local<Data>().value); });
1167     });
1168   });
1169
1170   fm.addTask([&]() {
1171     EXPECT_EQ(42, local<Data>().value);
1172
1173     local<Data>().value = 43;
1174
1175     fm.addTaskRemote([]() { EXPECT_EQ(43, local<Data>().value); });
1176   });
1177
1178   fm.addTask([]() {
1179     EXPECT_EQ(42, local<Data>().value);
1180     local<Data>().value = 43;
1181
1182     auto task = []() {
1183       EXPECT_EQ(43, local<Data>().value);
1184       local<Data>().value = 44;
1185     };
1186     std::vector<std::function<void()>> tasks{task};
1187     collectAny(tasks.begin(), tasks.end());
1188
1189     EXPECT_EQ(43, local<Data>().value);
1190   });
1191
1192   fm.loopUntilNoReady();
1193   EXPECT_FALSE(fm.hasTasks());
1194 }
1195
1196 TEST(FiberManager, fiberLocal) {
1197   struct SimpleData {
1198     int value{42};
1199   };
1200
1201   testFiberLocal<SimpleData>();
1202 }
1203
1204 TEST(FiberManager, fiberLocalHeap) {
1205   struct LargeData {
1206     char _[1024 * 1024];
1207     int value{42};
1208   };
1209
1210   testFiberLocal<LargeData>();
1211 }
1212
1213 TEST(FiberManager, fiberLocalDestructor) {
1214   struct CrazyData {
1215     size_t data{42};
1216
1217     ~CrazyData() {
1218       if (data == 41) {
1219         addTask([]() {
1220           EXPECT_EQ(42, local<CrazyData>().data);
1221           // Make sure we don't have infinite loop
1222           local<CrazyData>().data = 0;
1223         });
1224       }
1225     }
1226   };
1227
1228   FiberManager fm(
1229       LocalType<CrazyData>(), folly::make_unique<SimpleLoopController>());
1230
1231   fm.addTask([]() { local<CrazyData>().data = 41; });
1232
1233   fm.loopUntilNoReady();
1234   EXPECT_FALSE(fm.hasTasks());
1235 }
1236
1237 TEST(FiberManager, yieldTest) {
1238   FiberManager manager(folly::make_unique<SimpleLoopController>());
1239   auto& loopController =
1240       dynamic_cast<SimpleLoopController&>(manager.loopController());
1241
1242   bool checkRan = false;
1243
1244   manager.addTask([&]() {
1245     manager.yield();
1246     checkRan = true;
1247   });
1248
1249   loopController.loop([&]() {
1250     if (checkRan) {
1251       loopController.stop();
1252     }
1253   });
1254
1255   EXPECT_TRUE(checkRan);
1256 }
1257
1258 TEST(FiberManager, RequestContext) {
1259   FiberManager fm(folly::make_unique<SimpleLoopController>());
1260
1261   bool checkRun1 = false;
1262   bool checkRun2 = false;
1263   bool checkRun3 = false;
1264   bool checkRun4 = false;
1265   folly::fibers::Baton baton1;
1266   folly::fibers::Baton baton2;
1267   folly::fibers::Baton baton3;
1268   folly::fibers::Baton baton4;
1269
1270   folly::RequestContext::create();
1271   auto rcontext1 = folly::RequestContext::get();
1272   fm.addTask([&]() {
1273     EXPECT_EQ(rcontext1, folly::RequestContext::get());
1274     baton1.wait([&]() { EXPECT_EQ(rcontext1, folly::RequestContext::get()); });
1275     EXPECT_EQ(rcontext1, folly::RequestContext::get());
1276     runInMainContext(
1277         [&]() { EXPECT_EQ(rcontext1, folly::RequestContext::get()); });
1278     checkRun1 = true;
1279   });
1280
1281   folly::RequestContext::create();
1282   auto rcontext2 = folly::RequestContext::get();
1283   fm.addTaskRemote([&]() {
1284     EXPECT_EQ(rcontext2, folly::RequestContext::get());
1285     baton2.wait();
1286     EXPECT_EQ(rcontext2, folly::RequestContext::get());
1287     checkRun2 = true;
1288   });
1289
1290   folly::RequestContext::create();
1291   auto rcontext3 = folly::RequestContext::get();
1292   fm.addTaskFinally(
1293       [&]() {
1294         EXPECT_EQ(rcontext3, folly::RequestContext::get());
1295         baton3.wait();
1296         EXPECT_EQ(rcontext3, folly::RequestContext::get());
1297
1298         return folly::Unit();
1299       },
1300       [&](Try<folly::Unit>&& /* t */) {
1301         EXPECT_EQ(rcontext3, folly::RequestContext::get());
1302         checkRun3 = true;
1303       });
1304
1305   folly::RequestContext::setContext(nullptr);
1306   fm.addTask([&]() {
1307     folly::RequestContext::create();
1308     auto rcontext4 = folly::RequestContext::get();
1309     baton4.wait();
1310     EXPECT_EQ(rcontext4, folly::RequestContext::get());
1311     checkRun4 = true;
1312   });
1313
1314   folly::RequestContext::create();
1315   auto rcontext = folly::RequestContext::get();
1316
1317   fm.loopUntilNoReady();
1318   EXPECT_EQ(rcontext, folly::RequestContext::get());
1319
1320   baton1.post();
1321   EXPECT_EQ(rcontext, folly::RequestContext::get());
1322   fm.loopUntilNoReady();
1323   EXPECT_TRUE(checkRun1);
1324   EXPECT_EQ(rcontext, folly::RequestContext::get());
1325
1326   baton2.post();
1327   EXPECT_EQ(rcontext, folly::RequestContext::get());
1328   fm.loopUntilNoReady();
1329   EXPECT_TRUE(checkRun2);
1330   EXPECT_EQ(rcontext, folly::RequestContext::get());
1331
1332   baton3.post();
1333   EXPECT_EQ(rcontext, folly::RequestContext::get());
1334   fm.loopUntilNoReady();
1335   EXPECT_TRUE(checkRun3);
1336   EXPECT_EQ(rcontext, folly::RequestContext::get());
1337
1338   baton4.post();
1339   EXPECT_EQ(rcontext, folly::RequestContext::get());
1340   fm.loopUntilNoReady();
1341   EXPECT_TRUE(checkRun4);
1342   EXPECT_EQ(rcontext, folly::RequestContext::get());
1343 }
1344
1345 TEST(FiberManager, resizePeriodically) {
1346   FiberManager::Options opts;
1347   opts.fibersPoolResizePeriodMs = 300;
1348   opts.maxFibersPoolSize = 5;
1349
1350   FiberManager manager(folly::make_unique<EventBaseLoopController>(), opts);
1351
1352   folly::EventBase evb;
1353   dynamic_cast<EventBaseLoopController&>(manager.loopController())
1354       .attachEventBase(evb);
1355
1356   std::vector<Baton> batons(10);
1357
1358   size_t tasksRun = 0;
1359   for (size_t i = 0; i < 30; ++i) {
1360     manager.addTask([i, &batons, &tasksRun]() {
1361       ++tasksRun;
1362       // Keep some fibers active indefinitely
1363       if (i < batons.size()) {
1364         batons[i].wait();
1365       }
1366     });
1367   }
1368
1369   EXPECT_EQ(0, tasksRun);
1370   EXPECT_EQ(30, manager.fibersAllocated());
1371   EXPECT_EQ(0, manager.fibersPoolSize());
1372
1373   evb.loopOnce();
1374   EXPECT_EQ(30, tasksRun);
1375   EXPECT_EQ(30, manager.fibersAllocated());
1376   // Can go over maxFibersPoolSize, 10 of 30 fibers still active
1377   EXPECT_EQ(20, manager.fibersPoolSize());
1378
1379   std::this_thread::sleep_for(std::chrono::milliseconds(400));
1380   evb.loopOnce(); // no fibers active in this period
1381   EXPECT_EQ(30, manager.fibersAllocated());
1382   EXPECT_EQ(20, manager.fibersPoolSize());
1383
1384   std::this_thread::sleep_for(std::chrono::milliseconds(400));
1385   evb.loopOnce(); // should shrink fibers pool to maxFibersPoolSize
1386   EXPECT_EQ(15, manager.fibersAllocated());
1387   EXPECT_EQ(5, manager.fibersPoolSize());
1388
1389   for (size_t i = 0; i < batons.size(); ++i) {
1390     batons[i].post();
1391   }
1392   evb.loopOnce();
1393   EXPECT_EQ(15, manager.fibersAllocated());
1394   EXPECT_EQ(15, manager.fibersPoolSize());
1395
1396   std::this_thread::sleep_for(std::chrono::milliseconds(400));
1397   evb.loopOnce(); // 10 fibers active in last period
1398   EXPECT_EQ(10, manager.fibersAllocated());
1399   EXPECT_EQ(10, manager.fibersPoolSize());
1400
1401   std::this_thread::sleep_for(std::chrono::milliseconds(400));
1402   evb.loopOnce();
1403   EXPECT_EQ(5, manager.fibersAllocated());
1404   EXPECT_EQ(5, manager.fibersPoolSize());
1405 }
1406
1407 TEST(FiberManager, batonWaitTimeoutHandler) {
1408   FiberManager manager(folly::make_unique<EventBaseLoopController>());
1409
1410   folly::EventBase evb;
1411   dynamic_cast<EventBaseLoopController&>(manager.loopController())
1412       .attachEventBase(evb);
1413
1414   size_t fibersRun = 0;
1415   Baton baton;
1416   Baton::TimeoutHandler timeoutHandler;
1417
1418   manager.addTask([&]() {
1419     baton.wait(timeoutHandler);
1420     ++fibersRun;
1421   });
1422   manager.loopUntilNoReady();
1423
1424   EXPECT_FALSE(baton.try_wait());
1425   EXPECT_EQ(0, fibersRun);
1426
1427   timeoutHandler.scheduleTimeout(std::chrono::milliseconds(250));
1428   std::this_thread::sleep_for(std::chrono::milliseconds(500));
1429
1430   EXPECT_FALSE(baton.try_wait());
1431   EXPECT_EQ(0, fibersRun);
1432
1433   evb.loopOnce();
1434   manager.loopUntilNoReady();
1435
1436   EXPECT_EQ(1, fibersRun);
1437 }
1438
1439 TEST(FiberManager, batonWaitTimeoutMany) {
1440   FiberManager manager(folly::make_unique<EventBaseLoopController>());
1441
1442   folly::EventBase evb;
1443   dynamic_cast<EventBaseLoopController&>(manager.loopController())
1444       .attachEventBase(evb);
1445
1446   constexpr size_t kNumTimeoutTasks = 10000;
1447   size_t tasksCount = kNumTimeoutTasks;
1448
1449   // We add many tasks to hit timeout queue deallocation logic.
1450   for (size_t i = 0; i < kNumTimeoutTasks; ++i) {
1451     manager.addTask([&]() {
1452       Baton baton;
1453       Baton::TimeoutHandler timeoutHandler;
1454
1455       folly::fibers::addTask([&] {
1456         timeoutHandler.scheduleTimeout(std::chrono::milliseconds(1000));
1457       });
1458
1459       baton.wait(timeoutHandler);
1460       if (--tasksCount == 0) {
1461         evb.terminateLoopSoon();
1462       }
1463     });
1464   }
1465
1466   evb.loopForever();
1467 }
1468
1469 TEST(FiberManager, remoteFutureTest) {
1470   FiberManager fiberManager(folly::make_unique<SimpleLoopController>());
1471   auto& loopController =
1472       dynamic_cast<SimpleLoopController&>(fiberManager.loopController());
1473
1474   int testValue1 = 5;
1475   int testValue2 = 7;
1476   auto f1 = fiberManager.addTaskFuture([&]() { return testValue1; });
1477   auto f2 = fiberManager.addTaskRemoteFuture([&]() { return testValue2; });
1478   loopController.loop([&]() { loopController.stop(); });
1479   auto v1 = f1.get();
1480   auto v2 = f2.get();
1481
1482   EXPECT_EQ(v1, testValue1);
1483   EXPECT_EQ(v2, testValue2);
1484 }
1485
1486 // Test that a void function produes a Future<Unit>.
1487 TEST(FiberManager, remoteFutureVoidUnitTest) {
1488   FiberManager fiberManager(folly::make_unique<SimpleLoopController>());
1489   auto& loopController =
1490       dynamic_cast<SimpleLoopController&>(fiberManager.loopController());
1491
1492   bool ranLocal = false;
1493   folly::Future<folly::Unit> futureLocal =
1494       fiberManager.addTaskFuture([&]() { ranLocal = true; });
1495
1496   bool ranRemote = false;
1497   folly::Future<folly::Unit> futureRemote =
1498       fiberManager.addTaskRemoteFuture([&]() { ranRemote = true; });
1499
1500   loopController.loop([&]() { loopController.stop(); });
1501
1502   futureLocal.wait();
1503   ASSERT_TRUE(ranLocal);
1504
1505   futureRemote.wait();
1506   ASSERT_TRUE(ranRemote);
1507 }
1508
1509 TEST(FiberManager, nestedFiberManagers) {
1510   folly::EventBase outerEvb;
1511   folly::EventBase innerEvb;
1512
1513   getFiberManager(outerEvb).addTask([&]() {
1514     EXPECT_EQ(
1515         &getFiberManager(outerEvb), FiberManager::getFiberManagerUnsafe());
1516
1517     runInMainContext([&]() {
1518       getFiberManager(innerEvb).addTask([&]() {
1519         EXPECT_EQ(
1520             &getFiberManager(innerEvb), FiberManager::getFiberManagerUnsafe());
1521
1522         innerEvb.terminateLoopSoon();
1523       });
1524
1525       innerEvb.loopForever();
1526     });
1527
1528     EXPECT_EQ(
1529         &getFiberManager(outerEvb), FiberManager::getFiberManagerUnsafe());
1530
1531     outerEvb.terminateLoopSoon();
1532   });
1533
1534   outerEvb.loopForever();
1535 }
1536
1537 static size_t sNumAwaits;
1538
1539 void runBenchmark(size_t numAwaits, size_t toSend) {
1540   sNumAwaits = numAwaits;
1541
1542   FiberManager fiberManager(folly::make_unique<SimpleLoopController>());
1543   auto& loopController =
1544       dynamic_cast<SimpleLoopController&>(fiberManager.loopController());
1545
1546   std::queue<Promise<int>> pendingRequests;
1547   static const size_t maxOutstanding = 5;
1548
1549   auto loop = [&fiberManager, &loopController, &pendingRequests, &toSend]() {
1550     if (pendingRequests.size() == maxOutstanding || toSend == 0) {
1551       if (pendingRequests.empty()) {
1552         return;
1553       }
1554       pendingRequests.front().setValue(0);
1555       pendingRequests.pop();
1556     } else {
1557       fiberManager.addTask([&pendingRequests]() {
1558         for (size_t i = 0; i < sNumAwaits; ++i) {
1559           auto result = await([&pendingRequests](Promise<int> promise) {
1560             pendingRequests.push(std::move(promise));
1561           });
1562           DCHECK_EQ(result, 0);
1563         }
1564       });
1565
1566       if (--toSend == 0) {
1567         loopController.stop();
1568       }
1569     }
1570   };
1571
1572   loopController.loop(std::move(loop));
1573 }
1574
1575 BENCHMARK(FiberManagerBasicOneAwait, iters) {
1576   runBenchmark(1, iters);
1577 }
1578
1579 BENCHMARK(FiberManagerBasicFiveAwaits, iters) {
1580   runBenchmark(5, iters);
1581 }
1582
1583 BENCHMARK(FiberManagerCreateDestroy, iters) {
1584   for (size_t i = 0; i < iters; ++i) {
1585     folly::EventBase evb;
1586     auto& fm = folly::fibers::getFiberManager(evb);
1587     fm.addTask([]() {});
1588     evb.loop();
1589   }
1590 }
1591
1592 BENCHMARK(FiberManagerAllocateDeallocatePattern, iters) {
1593   static const size_t kNumAllocations = 10000;
1594
1595   FiberManager::Options opts;
1596   opts.maxFibersPoolSize = 0;
1597
1598   FiberManager fiberManager(folly::make_unique<SimpleLoopController>(), opts);
1599
1600   for (size_t iter = 0; iter < iters; ++iter) {
1601     EXPECT_EQ(0, fiberManager.fibersPoolSize());
1602
1603     size_t fibersRun = 0;
1604
1605     for (size_t i = 0; i < kNumAllocations; ++i) {
1606       fiberManager.addTask([&fibersRun] { ++fibersRun; });
1607       fiberManager.loopUntilNoReady();
1608     }
1609
1610     EXPECT_EQ(10000, fibersRun);
1611     EXPECT_EQ(0, fiberManager.fibersPoolSize());
1612   }
1613 }
1614
1615 BENCHMARK(FiberManagerAllocateLargeChunk, iters) {
1616   static const size_t kNumAllocations = 10000;
1617
1618   FiberManager::Options opts;
1619   opts.maxFibersPoolSize = 0;
1620
1621   FiberManager fiberManager(folly::make_unique<SimpleLoopController>(), opts);
1622
1623   for (size_t iter = 0; iter < iters; ++iter) {
1624     EXPECT_EQ(0, fiberManager.fibersPoolSize());
1625
1626     size_t fibersRun = 0;
1627
1628     for (size_t i = 0; i < kNumAllocations; ++i) {
1629       fiberManager.addTask([&fibersRun] { ++fibersRun; });
1630     }
1631
1632     fiberManager.loopUntilNoReady();
1633
1634     EXPECT_EQ(10000, fibersRun);
1635     EXPECT_EQ(0, fiberManager.fibersPoolSize());
1636   }
1637 }