folly/experimental/fibers/test/FibersTest.cpp

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