folly/futures/detail/FSM.h

   1 /*
   2  * Copyright 2014-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 #pragma once
  18
  19 #include <atomic>
  20 #include <mutex>
  21
  22 #include <folly/MicroSpinLock.h>
  23
  24 namespace folly {
  25 namespace futures {
  26 namespace detail {
  27
  28 /// Finite State Machine helper base class.
  29 /// Inherit from this.
  30 /// For best results, use an "enum class" for Enum.
  31 template <class Enum>
  32 class FSM {
  33  private:
  34   // I am not templatizing this because folly::MicroSpinLock needs to be
  35   // zero-initialized (or call init) which isn't generic enough for something
  36   // that behaves like std::mutex. :(
  37   using Mutex = folly::MicroSpinLock;
  38   Mutex mutex_ {0};
  39
  40   // This might not be necessary for all Enum types, e.g. anything
  41   // that is atomically updated in practice on this CPU and there's no risk
  42   // of returning a bogus state because of tearing.
  43   // An optimization would be to use a static conditional on the Enum type.
  44   std::atomic<Enum> state_;
  45
  46  public:
  47   explicit FSM(Enum startState) : state_(startState) {}
  48
  49   Enum getState() const noexcept {
  50     return state_.load(std::memory_order_acquire);
  51   }
  52
  53   /// Atomically do a state transition with accompanying action.
  54   /// The action will see the old state.
  55   /// @returns true on success, false and action unexecuted otherwise
  56   template <class F>
  57   bool updateState(Enum A, Enum B, F const& action) {
  58     if (!mutex_.try_lock()) {
  59       mutex_.lock();
  60     }
  61     if (state_.load(std::memory_order_acquire) != A) {
  62       mutex_.unlock();
  63       return false;
  64     }
  65     action();
  66     state_.store(B, std::memory_order_release);
  67     mutex_.unlock();
  68     return true;
  69   }
  70
  71   /// Atomically do a state transition with accompanying action. Then do the
  72   /// unprotected action without holding the lock. If the atomic transition
  73   /// fails, returns false and neither action was executed.
  74   ///
  75   /// This facilitates code like this:
  76   ///   bool done = false;
  77   ///   while (!done) {
  78   ///     switch (getState()) {
  79   ///     case State::Foo:
  80   ///       done = updateState(State::Foo, State::Bar,
  81   ///           [&]{ /* do protected stuff */ },
  82   ///           [&]{ /* do unprotected stuff */});
  83   ///       break;
  84   ///
  85   /// Which reads nicer than code like this:
  86   ///   while (true) {
  87   ///     switch (getState()) {
  88   ///     case State::Foo:
  89   ///       if (!updateState(State::Foo, State::Bar,
  90   ///           [&]{ /* do protected stuff */ })) {
  91   ///         continue;
  92   ///       }
  93   ///       /* do unprotected stuff */
  94   ///       return; // or otherwise break out of the loop
  95   ///
  96   /// The protected action will see the old state, and the unprotected action
  97   /// will see the new state.
  98   template <class F1, class F2>
  99   bool updateState(Enum A, Enum B,
 100                    F1 const& protectedAction, F2 const& unprotectedAction) {
 101     bool result = updateState(A, B, protectedAction);
 102     if (result) {
 103       unprotectedAction();
 104     }
 105     return result;
 106   }
 107 };
 108
 109 #define FSM_START(fsm) {\
 110     bool done = false; \
 111     while (!done) { auto state = fsm.getState(); switch (state) {
 112
 113 #define FSM_UPDATE2(fsm, b, protectedAction, unprotectedAction) \
 114     done = fsm.updateState(state, (b), (protectedAction), (unprotectedAction));
 115
 116 #define FSM_UPDATE(fsm, b, action) FSM_UPDATE2(fsm, (b), (action), []{})
 117
 118 #define FSM_CASE(fsm, a, b, action) \
 119   case (a): \
 120     FSM_UPDATE(fsm, (b), (action)); \
 121     break;
 122
 123 #define FSM_CASE2(fsm, a, b, protectedAction, unprotectedAction) \
 124   case (a): \
 125     FSM_UPDATE2(fsm, (b), (protectedAction), (unprotectedAction)); \
 126     break;
 127
 128 #define FSM_BREAK done = true; break;
 129 #define FSM_END }}}
 130
 131 } // namespace detail
 132 } // namespace futures
 133 } // namespace folly