stateMachine.hpp

Go to the documentation of this file.

#ifndef STATE_MACHINE_HPP
#define STATE_MACHINE_HPP

#include <cassert>
#include <chrono>
#include <cstdint>
#include <functional>
#include <iostream>
#include <limits>
#include <memory>
#include <mutex>
#include <queue>
#include <stdexcept>
#include <unordered_map>
#include <vector>

#include "bufArray.hpp"
#include "common.hpp"
#include "exceptions.hpp"
#include "spinlock.hpp"

template <class Identifier, class Packet> class StateMachine {
private:
    using ConnectionID = typename Identifier::ConnectionID;
    using Hasher = typename Identifier::Hasher;
    static constexpr auto timeoutIDInvalid = std::numeric_limits<uint32_t>::max();

public:
    /*
     * XXX -------------------------------------------- XXX
     *       Public declarations
     * XXX -------------------------------------------- XXX
     */

    struct State;
    class FunIface;

    using stateFun = std::function<void(State &, Packet *, FunIface &)>;

    using timeoutFun = std::function<void(State &, FunIface &)>;

    static constexpr auto StateIDInvalid = std::numeric_limits<StateID>::max();

    struct State {
    public:
        void *stateData;
        StateID state;
        uint32_t timeoutID;

        State() : stateData(nullptr), state(StateIDInvalid), timeoutID(timeoutIDInvalid){};
        State(StateID state, void *stateData)
            : stateData(stateData), state(state), timeoutID(timeoutIDInvalid){};
        State(const State &s)
            : stateData(s.stateData), state(s.state), timeoutID(s.timeoutID){};

        void set(const State &s) {
            stateData = s.stateData;
            state = s.state;
            timeoutID = s.timeoutID;
        }
    };

    /*
     * XXX -------------------------------------------- XXX
     *       Interface exposed to state functions
     * XXX -------------------------------------------- XXX
     */

    class FunIface {
    private:
        friend class StateMachine<Identifier, Packet>;

        StateMachine<Identifier, Packet> *sm;
        uint32_t pktIdx;
        BufArray<Packet> &pktsBA;
        ConnectionID &cID;
        State &state;
        bool sendPkt;
        bool immediateTransition;

        // Private -> nobody can misuse any FunIface objects
        FunIface(StateMachine<Identifier, Packet> *sm, uint32_t pktIdx,
            BufArray<Packet> &pktsBA, ConnectionID &cID, State &state)
            : sm(sm), pktIdx(pktIdx), pktsBA(pktsBA), cID(cID), state(state), sendPkt(true),
              immediateTransition(false){};

    public:
        ~FunIface() {
            if (!sendPkt) {
                pktsBA.markDropPkt(pktIdx);
            }

            // We are not in the endstate - checked by runPkt

            if (immediateTransition) {
                /*
                auto sfIt = sm->functions.find(state.state);
                if (sfIt == sm->functions.end()) {
                    D(
                        std::cout
                            << "FunIface::~FunIface() Didn't find a function for this state"
                            << std::endl;)
                    // Don't throw, we are in a destructor
                    // throw std::runtime_error("FunIface::~FunIface() No such function
                    // found");
                }

                D(std::cout << "Running Function" << std::endl;)
                (sfIt->second)(state, pktsBA[pktIdx], *this);
                */

                assert((sm->functions.size() - 1) >= state.state);
                auto fun = sm->functions[state.state];
                assert(fun != nullptr);

                D(std::cout << "Running Function" << std::endl;)
                fun(state, pktsBA[pktIdx], *this);

                if (state.state == sm->endStateID) {
                    D(std::cout << "Reached endStateID - deleting connection" << std::endl;)
                    sm->removeState(cID);
                }
            }
        }

        void freePkt() { sendPkt = false; }

        Packet *getPkt() {
            throw new std::runtime_error("StateMachine::FunIface::getPkt() not implemented");
        }

        void transition(StateID newState) { state.state = newState; }

        void transitionNow(StateID newState) {
            state.state = newState;
            this->immediateTransition = true;
        }

        void setTimeout(std::chrono::milliseconds timeout, timeoutFun fun) {
            std::chrono::time_point<std::chrono::steady_clock> now =
                std::chrono::steady_clock::now();
            std::chrono::time_point<std::chrono::steady_clock> then = now + timeout;

            struct Timeout t;
            t.time = then;
            t.timeoutID = sm->curTimeoutID++;

            state.timeoutID = t.timeoutID;

            sm->timeoutsQ.push(t);

            // This is an alternative to emplace
            // emplace should be better
            /*
            auto td = std::make_unique<struct TimeoutData>(cID, fun);
            sm->timeoutFunctions.insert(
                std::pair<uint32_t, std::unique_ptr<struct TimeoutData>>(
                    t.timeoutID, std::move(td)));
            */

            sm->timeoutFunctions.emplace(
                t.timeoutID, std::make_unique<struct TimeoutData>(cID, fun));
        }
    };

    class ConnectionPool {
    private:
        const unsigned int numBuckets = 8;
        const uint64_t bucketMask = 0b111;

        std::unordered_map<ConnectionID, State, Hasher> **newStates;
        SpinLockCLSize **newStatesLock;

    public:
        ConnectionPool() {
            newStates = reinterpret_cast<std::unordered_map<ConnectionID, State, Hasher> **>(
                malloc(sizeof(void *) * numBuckets));
            for (unsigned int i = 0; i < numBuckets; i++) {
                newStates[i] = new std::unordered_map<ConnectionID, State, Hasher>();
            }

            newStatesLock =
                reinterpret_cast<SpinLockCLSize **>(malloc(sizeof(void *) * numBuckets));
            for (unsigned int i = 0; i < numBuckets; i++) {
                newStatesLock[i] = new SpinLockCLSize();
            }
        };

        ~ConnectionPool() {
            for (unsigned int i = 0; i < numBuckets; i++) {
                delete (newStates[i]);
            }
            free(newStates);

            for (unsigned int i = 0; i < numBuckets; i++) {
                newStatesLock[i]->lock();
                delete (newStatesLock[i]);
            }
            free(newStatesLock);
        }

        void add(ConnectionID &cID, State &st) {
            Hasher hasher;
            uint64_t bucket = hasher(cID) & bucketMask;

            std::lock_guard<SpinLockCLSize> lock(*(newStatesLock[bucket]));
            auto insPair = std::pair<ConnectionID, State>(cID, st);
            newStates[bucket]->insert(insPair);
        };

        bool findAndErase(ConnectionID &cID, State *st) {
            Hasher hasher;
            uint64_t bucket = hasher(cID) & bucketMask;

            std::lock_guard<SpinLockCLSize> lock(*(newStatesLock[bucket]));
            auto ret = newStates[bucket]->find(cID);

            if (ret != newStates[bucket]->end()) {
                st->set(ret->second);
                newStates[bucket]->erase(ret);
                return true;
            } else {
                return false;
            }
        };

            // TODO I should really do this sooner than later
#if 0
        // If this is needed later, I will add it later
        void erase(ConnectionID &cID);
#endif
    };

private:
    /*
     * XXX -------------------------------------------- XXX
     *       Private attributes
     * XXX -------------------------------------------- XXX
     */

    // This is the heart of the state tracking
    // stateTable holds the link between connections and states
    std::unordered_map<ConnectionID, State, Hasher> stateTable;

    // This table specifies which function should be called for a packet
    // belonging to a connection in a specific state
    std::vector<stateFun> functions;

    // TODO: Since the identifier should not track anything, maybe we don't
    // even need a member -> all static functions
    Identifier identifier;

    // The state a newly received connection starts in
    // This is only useful, if listenToConnections is true
    StateID startStateID;

    // This function gets called on newly received connections
    // This is only useful, if listenToConnections is true
    std::function<void *(ConnectionID)> startStateFun;

    // If a connection reaches this state, it gets destryed
    StateID endStateID;

    // Callback to aquire new packets
    std::function<Packet *()> getPktCB;

    // Basically: Server mode or client mode
    bool listenToConnections;

    /*
     * XXX -------------------------------------------- XXX
     *       Timeout handling
     * XXX -------------------------------------------- XXX
     */

    // This represents the timeout tracking
    struct Timeout {
        // This is the, when the timeout ticks out
        std::chrono::time_point<std::chrono::steady_clock> time;

        // This uniquely identifies a single timeout
        uint32_t timeoutID;

        // Make sure, that the nearest timeout is the first one in timeoutsQ (see below)
        class Compare {
        public:
            bool operator()(struct Timeout a, struct Timeout b) { return a.time < b.time; };
        };
    };

    // This increments for every timeout -> provides a unique ID
    uint32_t curTimeoutID;

    // This contains all the timeouts
    std::priority_queue<struct Timeout, std::vector<struct Timeout>,
        typename Timeout::Compare>
        timeoutsQ;

    // This is used only to check if a timeout is valid, and given that, to
    // execute it
    struct TimeoutData {
        // This is used to identify, which connection the timeout belongs to
        // indexes the stateTable
        ConnectionID id;

        // This function is executed, if the timeout ticks out
        timeoutFun fun;

        // Trivial constructor
        TimeoutData(ConnectionID id, timeoutFun fun) : id(id), fun(fun){};
    };

    // If a timeoutID maps to a TimeoutData, this timeout is valid
    // if it maps to ::end(), then it is ignored.
    // The unique_ptr should make sure, that the map is not blown up by large ConnectionIDs
    // XXX The indirection assumes, that timeouts don't happen frequently, otherwise
    // it may be better to store them directly in the map? (not sure)
    std::unordered_map<uint32_t, std::unique_ptr<struct TimeoutData>> timeoutFunctions;

    /*
     * XXX -------------------------------------------- XXX
     *       Connection sharing
     * XXX -------------------------------------------- XXX
     */

    // These two members allow to open a connection on one core,
    // and receive subsequent packets on another
    static ConnectionPool connPoolStatic;
    ConnectionPool *connPool;

    /*
     * XXX -------------------------------------------- XXX
     *       Statistics
     * XXX -------------------------------------------- XXX
     */

    uint64_t stat_statesAdded = 0;
    uint64_t stat_statesClosed = 0;

    /*
     * XXX -------------------------------------------- XXX
     *       Private helper methods
     * XXX -------------------------------------------- XXX
     */

    auto findState(ConnectionID id) {
    findStateLoop:
        auto stateIt = stateTable.find(id);
        if (stateIt == stateTable.end()) {

            // Try to find state in the connection pool
            {
                State st;
                if (connPool->findAndErase(id, &st)) {
                    D(std::cout << "StateMachine::findState() found state in connPool"
                                << std::endl;)
                    stateTable.insert({id, st});

                    stat_statesAdded++;

                    goto findStateLoop;
                }
            }

            // Maybe accept the new connection
            if (listenToConnections) {
                // Add new state
                D(std::cout << "Adding new state" << std::endl;)
                D(std::cout << "ConnectionID: " << static_cast<std::string>(id) << std::endl;)

                // Create startState data object
                void *stateData = nullptr;
                if (startStateFun) {
                    stateData = startStateFun(id);
                }

                State s(startStateID, stateData);
                stateTable.insert({id, s});

                stat_statesAdded++;

                goto findStateLoop;
            }

        } else {
            D(std::cout << "State found" << std::endl;)
        }
        return stateIt;
    };

    void runPkt(BufArray<Packet> &pktsIn, unsigned int cur) {
        D(std::cout << std::endl << "StateMachine::runPkt() called" << std::endl;)

        try {
            // Retrieve the current packet
            Packet *pktIn = pktsIn[cur];

            // Try to identify the inbound packet
            ConnectionID identity = identifier.identify(pktIn);

            // Find a state/connection associated with this packet
            auto stateIt = findState(identity);

            if (stateIt == stateTable.end()) {
                // We don't want this packet
                D(std::cout << "StateMachine::runPkt() discarding packet" << std::endl;)
                D(std::cout << "ident of packet: " << static_cast<std::string>(identity)
                            << std::endl;)
                return;
            }

            // Invalidate any previous timeouts
            if (stateIt->second.timeoutID != timeoutIDInvalid) {
                this->timeoutFunctions.erase(stateIt->second.timeoutID);
                stateIt->second.timeoutID = timeoutIDInvalid;
            }

            // Try to retrieve an appropriate function
            /*
            auto sfIt = functions.find(stateIt->second.state);
            if (sfIt == functions.end()) {
                D(std::cout << "StateMachine::runPkt() Didn't find a function for this state"
                            << std::endl;)
                throw std::runtime_error("StateMachine::runPkt() No such function found");
            }
            */
            assert((functions.size() - 1) >= stateIt->second.state);
            auto fun = functions[stateIt->second.state];
            assert(fun != nullptr);

            // Create the custom function interface
            FunIface funIface(this, cur, pktsIn, identity, stateIt->second);

            // Run the function
            D(std::cout << "StateMachine::runPkt() Running Function" << std::endl;)
            D(std::cout << "StateMachine::runPkt() identity: "
                        << static_cast<std::string>(identity) << std::endl;)
            D(std::cout << "StateMachine::runPkt() hexdump of packet: " << std::endl;)
            D(hexdump(pktIn->getData(), pktIn->getDataLen());)
            //(sfIt->second)(stateIt->second, pktIn, funIface);
            fun(stateIt->second, pktIn, funIface);

            // Check if the endstate is reached
            if (stateIt->second.state == endStateID) {
                D(std::cout
                        << "StateMachine::runPkt() Reached endStateID - deleting connection"
                        << std::endl;)
                removeState(identity);

                stat_statesClosed++;
            }

            // At this point, the funIface is destroyed, and it is checked, if
            // any transitionNow calls were made.
            // XXX Therefore: DO NOT WRITE ANY CODE BELOW THIS COMMENT
            // (or at least give it a hard thought)

        } catch (PacketNotIdentified *e) {
            D(std::cout << "StateMachine::runPkt() Packet could not be identified"
                        << std::endl;);
            pktsIn.markDropPkt(cur);
        }
    }

public:
    /*
     * XXX -------------------------------------------- XXX
     *       Public interface
     * XXX -------------------------------------------- XXX
     */

    StateMachine()
        : startStateID(0), endStateID(StateIDInvalid), listenToConnections(false),
          curTimeoutID(0), connPool(&connPoolStatic){};

    ~StateMachine() {
        std::cout << "StateMachine stats:" << std::endl;
        std::cout << "stateTable.size() = " << stateTable.size() << std::endl;
        std::cout << "statesAdded  = " << stat_statesAdded << std::endl;
        std::cout << "statesClosed = " << stat_statesClosed << std::endl;
    }

    size_t getStateTableSize() { return stateTable.size(); };

    void registerFunction(StateID id, stateFun function) {
        // functions.insert({id, function});
        if ((static_cast<int>(functions.size()) - 1) < static_cast<int>(id)) {
            functions.resize(id + 1);
        }
        functions[id] = function;
    }

    void registerEndStateID(StateID endStateID) { this->endStateID = endStateID; }

    void registerStartStateID(
        StateID startStateID, std::function<void *(ConnectionID)> startStateFun) {
        this->startStateID = startStateID;
        this->startStateFun = startStateFun;
        listenToConnections = true;
    }

    void registerGetPktCB(std::function<Packet *()> fun) { getPktCB = fun; }

    void setConnectionPool(ConnectionPool *cp) { connPool = cp; }

    void removeState(ConnectionID id) { stateTable.erase(id); }

    void addState(ConnectionID id, State st, BufArray<Packet> &pktsIn) {

        /*
        auto sfIt = functions.find(st.state);
        if (sfIt == functions.end()) {
            throw std::runtime_error("StateMachine::addState() No such function found");
        }
        */

        assert((functions.size() - 1) >= st.state);
        auto fun = functions[st.state];
        assert(fun != nullptr);

        FunIface funIface(this, 0, pktsIn, id, st);

        D(std::cout << "StateMachine::addState() Running Function" << std::endl;)
        //(sfIt->second)(st, pktsIn[0], funIface);
        fun(st, pktsIn[0], funIface);

        if (st.state == endStateID) {
            D(std::cout << "StateMachine::addState() Reached endStateID - deleting connection"
                        << std::endl;)
            return;
        }

        D(std::cout << "StateMachine::addState() adding connection to newStates"
                    << std::endl;)
        D(std::cout << "StateMachine::addState() identity: " << static_cast<std::string>(id)
                    << std::endl;)
        connPool->add(id, st);
    }

    void runPktBatch(BufArray<Packet> &pktsIn) {
        uint32_t inCount = pktsIn.getTotalCount();

        D(std::cout << std::endl
                    << "StateMachine::runPktBatch() running incoming batch now, #Pkts: "
                    << inCount << std::endl;)

        // This loop handles the timeouts
        // It breaks, if there are no usable timeouts anymore
        // It will (usually) not run until timeoutsQ is empty
        while (!timeoutsQ.empty()) {
            // Get the next timeout
            auto timeoutElem = timeoutsQ.top();

            // If the current timeout is in the future -> break
            if (timeoutElem.time > std::chrono::steady_clock::now()) {
                break;
            }

            // Extract some info from the timeout
            uint32_t timeoutID = timeoutElem.timeoutID;
            auto timeoutDataIt = timeoutFunctions.find(timeoutID);

            // Check if the timeout is valid
            if (timeoutDataIt == timeoutFunctions.end()) {
                timeoutsQ.pop();
                continue;
            }

            // Prepare function call
            std::unique_ptr<struct TimeoutData> timeoutData =
                std::move(timeoutDataIt->second);
            auto stateIt = findState(timeoutData->id);
            assert(stateIt != stateTable.end());
            FunIface funIface(this, std::numeric_limits<uint32_t>::max(), pktsIn,
                timeoutData->id, stateIt->second);

            // Clear the timeoutID from the state
            stateIt->second.timeoutID = timeoutIDInvalid;

            // Call function
            timeoutData->fun(stateIt->second, funIface);

            // Check if we reached the end state
            if (stateIt->second.state == endStateID) {
                D(std::cout << "Reached endStateID - deleting connection" << std::endl;)
                removeState(timeoutData->id);
            }

            // Clean up
            timeoutsQ.pop();
            timeoutFunctions.erase(timeoutDataIt);
        }

        // Run all the usual incoming packets
        for (uint32_t i = 0; i < inCount; i++) {
            runPkt(pktsIn, i);
        }
    }
};

// Define static members of the state machine

// Don't try to understand the template stuff, it works...
template <class Identifier, class Packet>
typename StateMachine<Identifier, Packet>::ConnectionPool
    StateMachine<Identifier, Packet>::connPoolStatic;

/*
template <class Identifier, class Packet>
const uint64_t StateMachine<Identifier, Packet>::ConnectionPool::bucketMask;

template <class Identifier, class Packet>
const unsigned int StateMachine<Identifier, Packet>::ConnectionPool::numBuckets;
*/

#endif /* STATE_MACHINE_HPP */