folly/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/fibers/AddTasks.h>
  27 #include <folly/fibers/EventBaseLoopController.h>
  28 #include <folly/fibers/FiberManager.h>
  29 #include <folly/fibers/FiberManagerMap.h>
  30 #include <folly/fibers/GenericBaton.h>
  31 #include <folly/fibers/SimpleLoopController.h>
  32 #include <folly/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   {
1271     folly::RequestContextScopeGuard rctx;
1272     auto rcontext1 = folly::RequestContext::get();
1273     fm.addTask([&]() {
1274       EXPECT_EQ(rcontext1, folly::RequestContext::get());
1275       baton1.wait(
1276           [&]() { EXPECT_EQ(rcontext1, folly::RequestContext::get()); });
1277       EXPECT_EQ(rcontext1, folly::RequestContext::get());
1278       runInMainContext(
1279           [&]() { EXPECT_EQ(rcontext1, folly::RequestContext::get()); });
1280       checkRun1 = true;
1281     });
1282   }
1283   {
1284     folly::RequestContextScopeGuard rctx;
1285     auto rcontext2 = folly::RequestContext::get();
1286     fm.addTaskRemote([&]() {
1287       EXPECT_EQ(rcontext2, folly::RequestContext::get());
1288       baton2.wait();
1289       EXPECT_EQ(rcontext2, folly::RequestContext::get());
1290       checkRun2 = true;
1291     });
1292   }
1293   {
1294     folly::RequestContextScopeGuard rctx;
1295     auto rcontext3 = folly::RequestContext::get();
1296     fm.addTaskFinally(
1297         [&]() {
1298           EXPECT_EQ(rcontext3, folly::RequestContext::get());
1299           baton3.wait();
1300           EXPECT_EQ(rcontext3, folly::RequestContext::get());
1301
1302           return folly::Unit();
1303         },
1304         [&](Try<folly::Unit>&& /* t */) {
1305           EXPECT_EQ(rcontext3, folly::RequestContext::get());
1306           checkRun3 = true;
1307         });
1308   }
1309   {
1310     folly::RequestContext::setContext(nullptr);
1311     fm.addTask([&]() {
1312       folly::RequestContextScopeGuard rctx;
1313       auto rcontext4 = folly::RequestContext::get();
1314       baton4.wait();
1315       EXPECT_EQ(rcontext4, folly::RequestContext::get());
1316       checkRun4 = true;
1317     });
1318   }
1319   {
1320     folly::RequestContextScopeGuard rctx;
1321     auto rcontext = folly::RequestContext::get();
1322
1323     fm.loopUntilNoReady();
1324     EXPECT_EQ(rcontext, folly::RequestContext::get());
1325
1326     baton1.post();
1327     EXPECT_EQ(rcontext, folly::RequestContext::get());
1328     fm.loopUntilNoReady();
1329     EXPECT_TRUE(checkRun1);
1330     EXPECT_EQ(rcontext, folly::RequestContext::get());
1331
1332     baton2.post();
1333     EXPECT_EQ(rcontext, folly::RequestContext::get());
1334     fm.loopUntilNoReady();
1335     EXPECT_TRUE(checkRun2);
1336     EXPECT_EQ(rcontext, folly::RequestContext::get());
1337
1338     baton3.post();
1339     EXPECT_EQ(rcontext, folly::RequestContext::get());
1340     fm.loopUntilNoReady();
1341     EXPECT_TRUE(checkRun3);
1342     EXPECT_EQ(rcontext, folly::RequestContext::get());
1343
1344     baton4.post();
1345     EXPECT_EQ(rcontext, folly::RequestContext::get());
1346     fm.loopUntilNoReady();
1347     EXPECT_TRUE(checkRun4);
1348     EXPECT_EQ(rcontext, folly::RequestContext::get());
1349   }
1350 }
1351
1352 TEST(FiberManager, resizePeriodically) {
1353   FiberManager::Options opts;
1354   opts.fibersPoolResizePeriodMs = 300;
1355   opts.maxFibersPoolSize = 5;
1356
1357   FiberManager manager(folly::make_unique<EventBaseLoopController>(), opts);
1358
1359   folly::EventBase evb;
1360   dynamic_cast<EventBaseLoopController&>(manager.loopController())
1361       .attachEventBase(evb);
1362
1363   std::vector<Baton> batons(10);
1364
1365   size_t tasksRun = 0;
1366   for (size_t i = 0; i < 30; ++i) {
1367     manager.addTask([i, &batons, &tasksRun]() {
1368       ++tasksRun;
1369       // Keep some fibers active indefinitely
1370       if (i < batons.size()) {
1371         batons[i].wait();
1372       }
1373     });
1374   }
1375
1376   EXPECT_EQ(0, tasksRun);
1377   EXPECT_EQ(30, manager.fibersAllocated());
1378   EXPECT_EQ(0, manager.fibersPoolSize());
1379
1380   evb.loopOnce();
1381   EXPECT_EQ(30, tasksRun);
1382   EXPECT_EQ(30, manager.fibersAllocated());
1383   // Can go over maxFibersPoolSize, 10 of 30 fibers still active
1384   EXPECT_EQ(20, manager.fibersPoolSize());
1385
1386   std::this_thread::sleep_for(std::chrono::milliseconds(400));
1387   evb.loopOnce(); // no fibers active in this period
1388   EXPECT_EQ(30, manager.fibersAllocated());
1389   EXPECT_EQ(20, manager.fibersPoolSize());
1390
1391   std::this_thread::sleep_for(std::chrono::milliseconds(400));
1392   evb.loopOnce(); // should shrink fibers pool to maxFibersPoolSize
1393   EXPECT_EQ(15, manager.fibersAllocated());
1394   EXPECT_EQ(5, manager.fibersPoolSize());
1395
1396   for (size_t i = 0; i < batons.size(); ++i) {
1397     batons[i].post();
1398   }
1399   evb.loopOnce();
1400   EXPECT_EQ(15, manager.fibersAllocated());
1401   EXPECT_EQ(15, manager.fibersPoolSize());
1402
1403   std::this_thread::sleep_for(std::chrono::milliseconds(400));
1404   evb.loopOnce(); // 10 fibers active in last period
1405   EXPECT_EQ(10, manager.fibersAllocated());
1406   EXPECT_EQ(10, manager.fibersPoolSize());
1407
1408   std::this_thread::sleep_for(std::chrono::milliseconds(400));
1409   evb.loopOnce();
1410   EXPECT_EQ(5, manager.fibersAllocated());
1411   EXPECT_EQ(5, manager.fibersPoolSize());
1412 }
1413
1414 TEST(FiberManager, batonWaitTimeoutHandler) {
1415   FiberManager manager(folly::make_unique<EventBaseLoopController>());
1416
1417   folly::EventBase evb;
1418   dynamic_cast<EventBaseLoopController&>(manager.loopController())
1419       .attachEventBase(evb);
1420
1421   size_t fibersRun = 0;
1422   Baton baton;
1423   Baton::TimeoutHandler timeoutHandler;
1424
1425   manager.addTask([&]() {
1426     baton.wait(timeoutHandler);
1427     ++fibersRun;
1428   });
1429   manager.loopUntilNoReady();
1430
1431   EXPECT_FALSE(baton.try_wait());
1432   EXPECT_EQ(0, fibersRun);
1433
1434   timeoutHandler.scheduleTimeout(std::chrono::milliseconds(250));
1435   std::this_thread::sleep_for(std::chrono::milliseconds(500));
1436
1437   EXPECT_FALSE(baton.try_wait());
1438   EXPECT_EQ(0, fibersRun);
1439
1440   evb.loopOnce();
1441   manager.loopUntilNoReady();
1442
1443   EXPECT_EQ(1, fibersRun);
1444 }
1445
1446 TEST(FiberManager, batonWaitTimeoutMany) {
1447   FiberManager manager(folly::make_unique<EventBaseLoopController>());
1448
1449   folly::EventBase evb;
1450   dynamic_cast<EventBaseLoopController&>(manager.loopController())
1451       .attachEventBase(evb);
1452
1453   constexpr size_t kNumTimeoutTasks = 10000;
1454   size_t tasksCount = kNumTimeoutTasks;
1455
1456   // We add many tasks to hit timeout queue deallocation logic.
1457   for (size_t i = 0; i < kNumTimeoutTasks; ++i) {
1458     manager.addTask([&]() {
1459       Baton baton;
1460       Baton::TimeoutHandler timeoutHandler;
1461
1462       folly::fibers::addTask([&] {
1463         timeoutHandler.scheduleTimeout(std::chrono::milliseconds(1000));
1464       });
1465
1466       baton.wait(timeoutHandler);
1467       if (--tasksCount == 0) {
1468         evb.terminateLoopSoon();
1469       }
1470     });
1471   }
1472
1473   evb.loopForever();
1474 }
1475
1476 TEST(FiberManager, remoteFutureTest) {
1477   FiberManager fiberManager(folly::make_unique<SimpleLoopController>());
1478   auto& loopController =
1479       dynamic_cast<SimpleLoopController&>(fiberManager.loopController());
1480
1481   int testValue1 = 5;
1482   int testValue2 = 7;
1483   auto f1 = fiberManager.addTaskFuture([&]() { return testValue1; });
1484   auto f2 = fiberManager.addTaskRemoteFuture([&]() { return testValue2; });
1485   loopController.loop([&]() { loopController.stop(); });
1486   auto v1 = f1.get();
1487   auto v2 = f2.get();
1488
1489   EXPECT_EQ(v1, testValue1);
1490   EXPECT_EQ(v2, testValue2);
1491 }
1492
1493 // Test that a void function produes a Future<Unit>.
1494 TEST(FiberManager, remoteFutureVoidUnitTest) {
1495   FiberManager fiberManager(folly::make_unique<SimpleLoopController>());
1496   auto& loopController =
1497       dynamic_cast<SimpleLoopController&>(fiberManager.loopController());
1498
1499   bool ranLocal = false;
1500   folly::Future<folly::Unit> futureLocal =
1501       fiberManager.addTaskFuture([&]() { ranLocal = true; });
1502
1503   bool ranRemote = false;
1504   folly::Future<folly::Unit> futureRemote =
1505       fiberManager.addTaskRemoteFuture([&]() { ranRemote = true; });
1506
1507   loopController.loop([&]() { loopController.stop(); });
1508
1509   futureLocal.wait();
1510   ASSERT_TRUE(ranLocal);
1511
1512   futureRemote.wait();
1513   ASSERT_TRUE(ranRemote);
1514 }
1515
1516 TEST(FiberManager, nestedFiberManagers) {
1517   folly::EventBase outerEvb;
1518   folly::EventBase innerEvb;
1519
1520   getFiberManager(outerEvb).addTask([&]() {
1521     EXPECT_EQ(
1522         &getFiberManager(outerEvb), FiberManager::getFiberManagerUnsafe());
1523
1524     runInMainContext([&]() {
1525       getFiberManager(innerEvb).addTask([&]() {
1526         EXPECT_EQ(
1527             &getFiberManager(innerEvb), FiberManager::getFiberManagerUnsafe());
1528
1529         innerEvb.terminateLoopSoon();
1530       });
1531
1532       innerEvb.loopForever();
1533     });
1534
1535     EXPECT_EQ(
1536         &getFiberManager(outerEvb), FiberManager::getFiberManagerUnsafe());
1537
1538     outerEvb.terminateLoopSoon();
1539   });
1540
1541   outerEvb.loopForever();
1542 }
1543
1544 static size_t sNumAwaits;
1545
1546 void runBenchmark(size_t numAwaits, size_t toSend) {
1547   sNumAwaits = numAwaits;
1548
1549   FiberManager fiberManager(folly::make_unique<SimpleLoopController>());
1550   auto& loopController =
1551       dynamic_cast<SimpleLoopController&>(fiberManager.loopController());
1552
1553   std::queue<Promise<int>> pendingRequests;
1554   static const size_t maxOutstanding = 5;
1555
1556   auto loop = [&fiberManager, &loopController, &pendingRequests, &toSend]() {
1557     if (pendingRequests.size() == maxOutstanding || toSend == 0) {
1558       if (pendingRequests.empty()) {
1559         return;
1560       }
1561       pendingRequests.front().setValue(0);
1562       pendingRequests.pop();
1563     } else {
1564       fiberManager.addTask([&pendingRequests]() {
1565         for (size_t i = 0; i < sNumAwaits; ++i) {
1566           auto result = await([&pendingRequests](Promise<int> promise) {
1567             pendingRequests.push(std::move(promise));
1568           });
1569           DCHECK_EQ(result, 0);
1570         }
1571       });
1572
1573       if (--toSend == 0) {
1574         loopController.stop();
1575       }
1576     }
1577   };
1578
1579   loopController.loop(std::move(loop));
1580 }
1581
1582 BENCHMARK(FiberManagerBasicOneAwait, iters) {
1583   runBenchmark(1, iters);
1584 }
1585
1586 BENCHMARK(FiberManagerBasicFiveAwaits, iters) {
1587   runBenchmark(5, iters);
1588 }
1589
1590 BENCHMARK(FiberManagerCreateDestroy, iters) {
1591   for (size_t i = 0; i < iters; ++i) {
1592     folly::EventBase evb;
1593     auto& fm = folly::fibers::getFiberManager(evb);
1594     fm.addTask([]() {});
1595     evb.loop();
1596   }
1597 }
1598
1599 BENCHMARK(FiberManagerAllocateDeallocatePattern, iters) {
1600   static const size_t kNumAllocations = 10000;
1601
1602   FiberManager::Options opts;
1603   opts.maxFibersPoolSize = 0;
1604
1605   FiberManager fiberManager(folly::make_unique<SimpleLoopController>(), opts);
1606
1607   for (size_t iter = 0; iter < iters; ++iter) {
1608     EXPECT_EQ(0, fiberManager.fibersPoolSize());
1609
1610     size_t fibersRun = 0;
1611
1612     for (size_t i = 0; i < kNumAllocations; ++i) {
1613       fiberManager.addTask([&fibersRun] { ++fibersRun; });
1614       fiberManager.loopUntilNoReady();
1615     }
1616
1617     EXPECT_EQ(10000, fibersRun);
1618     EXPECT_EQ(0, fiberManager.fibersPoolSize());
1619   }
1620 }
1621
1622 BENCHMARK(FiberManagerAllocateLargeChunk, iters) {
1623   static const size_t kNumAllocations = 10000;
1624
1625   FiberManager::Options opts;
1626   opts.maxFibersPoolSize = 0;
1627
1628   FiberManager fiberManager(folly::make_unique<SimpleLoopController>(), opts);
1629
1630   for (size_t iter = 0; iter < iters; ++iter) {
1631     EXPECT_EQ(0, fiberManager.fibersPoolSize());
1632
1633     size_t fibersRun = 0;
1634
1635     for (size_t i = 0; i < kNumAllocations; ++i) {
1636       fiberManager.addTask([&fibersRun] { ++fibersRun; });
1637     }
1638
1639     fiberManager.loopUntilNoReady();
1640
1641     EXPECT_EQ(10000, fibersRun);
1642     EXPECT_EQ(0, fiberManager.fibersPoolSize());
1643   }
1644 }