folly/io/async/test/RequestContextTest.cpp

   1 /*
   2  * Copyright 2013-present 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
  17 #include <thread>
  18
  19 #include <folly/Memory.h>
  20 #include <folly/io/async/EventBase.h>
  21 #include <folly/io/async/Request.h>
  22 #include <folly/portability/GTest.h>
  23
  24 using namespace folly;
  25
  26 class TestData : public RequestData {
  27  public:
  28   explicit TestData(int data) : data_(data) {}
  29   ~TestData() override {}
  30
  31   bool hasCallback() override {
  32     return true;
  33   }
  34
  35   void onSet() override {
  36     set_++;
  37   }
  38
  39   void onUnset() override {
  40     unset_++;
  41   }
  42
  43   int set_ = 0, unset_ = 0;
  44   int data_;
  45 };
  46
  47 TEST(RequestContext, SimpleTest) {
  48   EventBase base;
  49
  50
  51   // There should always be a default context with get()
  52   EXPECT_TRUE(RequestContext::get() != nullptr);
  53
  54
  55   // but not with saveContext()
  56   EXPECT_EQ(RequestContext::saveContext(), nullptr);
  57   RequestContext::create();
  58   EXPECT_NE(RequestContext::saveContext(), nullptr);
  59   RequestContext::create();
  60   EXPECT_NE(RequestContext::saveContext(), nullptr);
  61
  62   EXPECT_EQ(nullptr, RequestContext::get()->getContextData("test"));
  63
  64   RequestContext::get()->setContextData("test", std::make_unique<TestData>(10));
  65   base.runInEventBaseThread([&](){
  66       EXPECT_TRUE(RequestContext::get() != nullptr);
  67       auto data = dynamic_cast<TestData*>(
  68         RequestContext::get()->getContextData("test"))->data_;
  69       EXPECT_EQ(10, data);
  70       base.terminateLoopSoon();
  71     });
  72   auto th = std::thread([&](){
  73       base.loopForever();
  74   });
  75   th.join();
  76   EXPECT_TRUE(RequestContext::get() != nullptr);
  77   auto a = dynamic_cast<TestData*>(
  78     RequestContext::get()->getContextData("test"));
  79   auto data = a->data_;
  80   EXPECT_EQ(10, data);
  81
  82   RequestContext::setContext(std::shared_ptr<RequestContext>());
  83   // There should always be a default context
  84   EXPECT_TRUE(nullptr != RequestContext::get());
  85 }
  86
  87 TEST(RequestContext, setIfAbsentTest) {
  88   EXPECT_TRUE(RequestContext::get() != nullptr);
  89
  90   RequestContext::get()->setContextData("test", std::make_unique<TestData>(10));
  91   EXPECT_FALSE(RequestContext::get()->setContextDataIfAbsent(
  92       "test", std::make_unique<TestData>(20)));
  93   EXPECT_EQ(10,
  94             dynamic_cast<TestData*>(
  95                 RequestContext::get()->getContextData("test"))->data_);
  96
  97   EXPECT_TRUE(RequestContext::get()->setContextDataIfAbsent(
  98       "test2", std::make_unique<TestData>(20)));
  99   EXPECT_EQ(20,
 100             dynamic_cast<TestData*>(
 101                 RequestContext::get()->getContextData("test2"))->data_);
 102
 103   RequestContext::setContext(std::shared_ptr<RequestContext>());
 104   EXPECT_TRUE(nullptr != RequestContext::get());
 105 }
 106
 107 TEST(RequestContext, testSetUnset) {
 108   RequestContext::create();
 109   auto ctx1 = RequestContext::saveContext();
 110   ctx1->setContextData("test", std::make_unique<TestData>(10));
 111   auto testData1 = dynamic_cast<TestData*>(ctx1->getContextData("test"));
 112
 113   // Override RequestContext
 114   RequestContext::create();
 115   auto ctx2 = RequestContext::saveContext();
 116   ctx2->setContextData("test", std::make_unique<TestData>(20));
 117   auto testData2 = dynamic_cast<TestData*>(ctx2->getContextData("test"));
 118
 119   // Check ctx1->onUnset was called
 120   EXPECT_EQ(0, testData1->set_);
 121   EXPECT_EQ(1, testData1->unset_);
 122
 123   RequestContext::setContext(ctx1);
 124   EXPECT_EQ(1, testData1->set_);
 125   EXPECT_EQ(1, testData1->unset_);
 126   EXPECT_EQ(0, testData2->set_);
 127   EXPECT_EQ(1, testData2->unset_);
 128
 129   RequestContext::setContext(ctx2);
 130   EXPECT_EQ(1, testData1->set_);
 131   EXPECT_EQ(2, testData1->unset_);
 132   EXPECT_EQ(1, testData2->set_);
 133   EXPECT_EQ(1, testData2->unset_);
 134 }
 135
 136 TEST(RequestContext, deadlockTest) {
 137   class DeadlockTestData : public RequestData {
 138    public:
 139     explicit DeadlockTestData(const std::string& val) : val_(val) {}
 140
 141     ~DeadlockTestData() override {
 142       RequestContext::get()->setContextData(
 143           val_, std::make_unique<TestData>(1));
 144     }
 145
 146     bool hasCallback() override {
 147       return false;
 148     }
 149
 150     std::string val_;
 151   };
 152
 153   RequestContext::get()->setContextData(
 154       "test", std::make_unique<DeadlockTestData>("test2"));
 155   RequestContext::get()->clearContextData("test");
 156 }