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