fiber->state_ == Fiber::NOT_STARTED ||
fiber->state_ == Fiber::READY_TO_RUN);
currentFiber_ = fiber;
- fiber->rcontext_ = RequestContext::setContext(std::move(fiber->rcontext_));
if (observer_) {
observer_->starting(reinterpret_cast<uintptr_t>(fiber));
}
observer_->stopped(reinterpret_cast<uintptr_t>(fiber));
}
currentFiber_ = nullptr;
- fiber->rcontext_ = RequestContext::setContext(std::move(fiber->rcontext_));
} else if (fiber->state_ == Fiber::INVALID) {
assert(fibersActive_ > 0);
--fibersActive_;
observer_->stopped(reinterpret_cast<uintptr_t>(fiber));
}
currentFiber_ = nullptr;
- fiber->rcontext_ = RequestContext::setContext(std::move(fiber->rcontext_));
fiber->localData_.reset();
fiber->rcontext_.reset();
observer_->stopped(reinterpret_cast<uintptr_t>(fiber));
}
currentFiber_ = nullptr;
- fiber->rcontext_ = RequestContext::setContext(std::move(fiber->rcontext_));
fiber->state_ = Fiber::READY_TO_RUN;
yieldedFibers_.push_back(*fiber);
}
auto originalFiberManager = this;
std::swap(currentFiberManager_, originalFiberManager);
+ RequestContext::Provider oldRequestContextProvider;
+ auto newRequestContextProvider =
+ [this, &oldRequestContextProvider]() -> std::shared_ptr<RequestContext>& {
+ return currentFiber_ ? currentFiber_->rcontext_
+ : oldRequestContextProvider();
+ };
+ oldRequestContextProvider = RequestContext::setRequestContextProvider(
+ std::ref(newRequestContextProvider));
+
SCOPE_EXIT {
isLoopScheduled_ = false;
+ // Restore RequestContext provider before call to ensureLoopScheduled()
+ RequestContext::setRequestContextProvider(
+ std::move(oldRequestContextProvider));
if (!readyFibers_.empty()) {
ensureLoopScheduled();
}
#include <folly/fibers/SimpleLoopController.h>
#include <folly/fibers/TimedMutex.h>
#include <folly/fibers/WhenN.h>
+#include <folly/io/async/Request.h>
#include <folly/io/async/ScopedEventBaseThread.h>
#include <folly/portability/GTest.h>
EXPECT_FALSE(fm.hasTasks());
}
+TEST(FiberManager, fiberRequestContext) {
+ folly::EventBase evb;
+ FiberManager fm(std::make_unique<EventBaseLoopController>());
+ dynamic_cast<EventBaseLoopController&>(fm.loopController())
+ .attachEventBase(evb);
+
+ struct TestContext : public folly::RequestData {
+ explicit TestContext(std::string s) : data(std::move(s)) {}
+ std::string data;
+ };
+
+ class AfterFibersCallback : public folly::EventBase::LoopCallback {
+ public:
+ AfterFibersCallback(
+ folly::EventBase& evb,
+ const bool& fibersDone,
+ folly::Function<void()> afterFibersFunc)
+ : evb_(evb),
+ fibersDone_(fibersDone),
+ afterFibersFunc_(std::move(afterFibersFunc)) {}
+
+ void runLoopCallback() noexcept override {
+ if (fibersDone_) {
+ afterFibersFunc_();
+ delete this;
+ } else {
+ evb_.runInLoop(this);
+ }
+ }
+
+ private:
+ folly::EventBase& evb_;
+ const bool& fibersDone_;
+ folly::Function<void()> afterFibersFunc_;
+ };
+
+ bool fibersDone = false;
+ size_t tasksRun = 0;
+ evb.runInEventBaseThread([&evb, &fm, &tasksRun, &fibersDone]() {
+ ++tasksRun;
+ auto* const evbCtx = folly::RequestContext::get();
+ EXPECT_NE(nullptr, evbCtx);
+ EXPECT_EQ(nullptr, evbCtx->getContextData("key"));
+ evbCtx->setContextData("key", std::make_unique<TestContext>("evb_value"));
+
+ // This callback allows us to check that FiberManager has restored the
+ // RequestContext provider as expected after a fiber loop.
+ auto* afterFibersCallback =
+ new AfterFibersCallback(evb, fibersDone, [&tasksRun, evbCtx]() {
+ ++tasksRun;
+ EXPECT_EQ(evbCtx, folly::RequestContext::get());
+ EXPECT_EQ(
+ "evb_value",
+ dynamic_cast<TestContext*>(evbCtx->getContextData("key"))->data);
+ });
+ evb.runInLoop(afterFibersCallback);
+
+ // Launching a fiber allows us to hit FiberManager RequestContext
+ // setup/teardown logic.
+ fm.addTask([&evb, &tasksRun, &fibersDone, evbCtx]() {
+ ++tasksRun;
+
+ // Initially, fiber starts with same RequestContext as its parent task.
+ EXPECT_EQ(evbCtx, folly::RequestContext::get());
+ EXPECT_NE(nullptr, evbCtx->getContextData("key"));
+ EXPECT_EQ(
+ "evb_value",
+ dynamic_cast<TestContext*>(evbCtx->getContextData("key"))->data);
+
+ // Create a new RequestContext for this fiber so we can distinguish from
+ // RequestContext first EventBase callback started with.
+ folly::RequestContext::create();
+ auto* const fiberCtx = folly::RequestContext::get();
+ EXPECT_NE(nullptr, fiberCtx);
+ EXPECT_EQ(nullptr, fiberCtx->getContextData("key"));
+ fiberCtx->setContextData(
+ "key", std::make_unique<TestContext>("fiber_value"));
+
+ // Task launched from within fiber should share current fiber's
+ // RequestContext
+ evb.runInEventBaseThread([&tasksRun, fiberCtx]() {
+ ++tasksRun;
+ auto* const evbCtx2 = folly::RequestContext::get();
+ EXPECT_EQ(fiberCtx, evbCtx2);
+ EXPECT_NE(nullptr, evbCtx2->getContextData("key"));
+ EXPECT_EQ(
+ "fiber_value",
+ dynamic_cast<TestContext*>(evbCtx2->getContextData("key"))->data);
+ });
+
+ fibersDone = true;
+ });
+ });
+
+ evb.loop();
+
+ EXPECT_EQ(4, tasksRun);
+ EXPECT_TRUE(fibersDone);
+ EXPECT_FALSE(fm.hasTasks());
+}
+
TEST(FiberManager, yieldTest) {
FiberManager manager(std::make_unique<SimpleLoopController>());
auto& loopController =
* See the License for the specific language governing permissions and
* limitations under the License.
*/
-
#include <folly/io/async/Request.h>
-#include <folly/tracing/StaticTracepoint.h>
+
+#include <algorithm>
+#include <stdexcept>
+#include <utility>
#include <glog/logging.h>
#include <folly/MapUtil.h>
#include <folly/SingletonThreadLocal.h>
+#include <folly/tracing/StaticTracepoint.h>
namespace folly {
return ctx;
}
-std::shared_ptr<RequestContext>& RequestContext::getStaticContext() {
- using SingletonT = SingletonThreadLocal<std::shared_ptr<RequestContext>>;
- static SingletonT singleton;
+RequestContext::Provider& RequestContext::requestContextProvider() {
+ class DefaultProvider {
+ public:
+ constexpr DefaultProvider() = default;
+ DefaultProvider(const DefaultProvider&) = delete;
+ DefaultProvider& operator=(const DefaultProvider&) = delete;
+ DefaultProvider(DefaultProvider&&) = default;
+ DefaultProvider& operator=(DefaultProvider&&) = default;
+
+ std::shared_ptr<RequestContext>& operator()() {
+ return context;
+ }
+
+ private:
+ std::shared_ptr<RequestContext> context;
+ };
- return singleton.get();
+ static SingletonThreadLocal<Provider> providerSingleton(
+ []() { return new Provider(DefaultProvider()); });
+ return providerSingleton.get();
+}
+
+std::shared_ptr<RequestContext>& RequestContext::getStaticContext() {
+ auto& provider = requestContextProvider();
+ return provider();
}
RequestContext* RequestContext::get() {
- auto context = getStaticContext();
+ auto& context = getStaticContext();
if (!context) {
static RequestContext defaultContext;
return std::addressof(defaultContext);
}
return context.get();
}
+
+RequestContext::Provider RequestContext::setRequestContextProvider(
+ RequestContext::Provider newProvider) {
+ if (!newProvider) {
+ throw std::runtime_error("RequestContext provider must be non-empty");
+ }
+
+ auto& provider = requestContextProvider();
+ std::swap(provider, newProvider);
+ return newProvider;
+}
}
#include <map>
#include <memory>
+#include <folly/Function.h>
#include <folly/SharedMutex.h>
#include <folly/Synchronized.h>
// copied between threads.
class RequestContext {
public:
+ using Provider = folly::Function<std::shared_ptr<RequestContext>&()>;
+
// Create a unique request context for this request.
// It will be passed between queues / threads (where implemented),
// so it should be valid for the lifetime of the request.
return getStaticContext();
}
+ // This API allows one to override the default behavior of getStaticContext()
+ // by providing a custom RequestContext provider. The old provider is
+ // returned, and the user must restore the old provider via a subsequent call
+ // to setRequestContextProvider() once the new provider is no longer needed.
+ //
+ // Using custom RequestContext providers can be more efficient than having to
+ // setContext() whenever context must be switched. This is especially true in
+ // applications that do not actually use RequestContext, but where library
+ // code must still support RequestContext for other use cases. See
+ // FiberManager for an example of how a custom RequestContext provider can
+ // reduce calls to setContext().
+ static Provider setRequestContextProvider(Provider f);
+
private:
static std::shared_ptr<RequestContext>& getStaticContext();
+ static Provider& requestContextProvider();
using Data = std::map<std::string, std::unique_ptr<RequestData>>;
folly::Synchronized<Data, folly::SharedMutex> data_;
EXPECT_TRUE(nullptr != RequestContext::get());
}
+TEST(RequestContext, nonDefaultContextsAreThreadLocal) {
+ RequestContext* ctx1 = nullptr;
+ RequestContext* ctx2 = nullptr;
+
+ std::vector<std::thread> ts;
+ for (size_t i = 0; i < 2; ++i) {
+ auto*& ctx = (i == 0 ? ctx1 : ctx2);
+ ts.emplace_back([&ctx]() {
+ RequestContext::create();
+ ctx = RequestContext::get();
+ });
+ }
+ for (auto& t : ts) {
+ t.join();
+ }
+
+ EXPECT_NE(nullptr, ctx1);
+ EXPECT_NE(nullptr, ctx2);
+ EXPECT_NE(ctx1, ctx2);
+}
+
+TEST(RequestContext, customRequestContextProvider) {
+ auto customContext = std::make_shared<RequestContext>();
+ auto customProvider = [&customContext]() -> std::shared_ptr<RequestContext>& {
+ return customContext;
+ };
+
+ auto* const originalContext = RequestContext::get();
+ EXPECT_NE(nullptr, originalContext);
+
+ // Install new RequestContext provider
+ auto originalProvider =
+ RequestContext::setRequestContextProvider(std::move(customProvider));
+
+ auto* const newContext = RequestContext::get();
+ EXPECT_EQ(customContext.get(), newContext);
+ EXPECT_NE(originalContext, newContext);
+
+ // Restore original RequestContext provider
+ RequestContext::setRequestContextProvider(std::move(originalProvider));
+ EXPECT_EQ(originalContext, RequestContext::get());
+}
+
TEST(RequestContext, setIfAbsentTest) {
EXPECT_TRUE(RequestContext::get() != nullptr);
projects / folly.git / commitdiff