folly/io/async/test/RequestContextTest.cpp

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