lokinet/llarp/dht/txholder.hpp

#ifndef LLARP_DHT_TXHOLDER
#define LLARP_DHT_TXHOLDER

#include <dht/tx.hpp>
#include <dht/txowner.hpp>
#include <util/time.hpp>
#include <util/status.hpp>

#include <memory>
#include <unordered_map>

namespace llarp
{
  namespace dht
  {
    template < typename K, typename V, typename K_Hash >
    struct TXHolder
    {
      using TXPtr = std::unique_ptr< TX< K, V > >;
      // tx who are waiting for a reply for each key
      std::unordered_multimap< K, TXOwner, K_Hash > waiting;
      // tx timesouts by key
      std::unordered_map< K, llarp_time_t, K_Hash > timeouts;
      // maps remote peer with tx to handle reply from them
      std::unordered_map< TXOwner, TXPtr, TXOwner::Hash > tx;

      const TX< K, V >*
      GetPendingLookupFrom(const TXOwner& owner) const;

      util::StatusObject
      ExtractStatus() const
      {
        util::StatusObject obj{};
        std::vector< util::StatusObject > txObjs, timeoutsObjs, waitingObjs;
        std::transform(tx.begin(), tx.end(), std::back_inserter(txObjs),
                       [](const auto& item) -> util::StatusObject {
                         return util::StatusObject{
                             {"owner", item.first.ExtractStatus()},
                             {"tx", item.second->ExtractStatus()}};
                       });
        obj["tx"] = txObjs;
        std::transform(
            timeouts.begin(), timeouts.end(), std::back_inserter(timeoutsObjs),
            [](const auto& item) -> util::StatusObject {
              return util::StatusObject{{"time", to_json(item.second)},
                                        {"target", item.first.ExtractStatus()}};
            });
        obj["timeouts"] = timeoutsObjs;
        std::transform(waiting.begin(), waiting.end(),
                       std::back_inserter(waitingObjs),
                       [](const auto& item) -> util::StatusObject {
                         return util::StatusObject{
                             {"target", item.first.ExtractStatus()},
                             {"whoasked", item.second.ExtractStatus()}};
                       });
        obj["waiting"] = waitingObjs;
        return obj;
      }

      bool
      HasLookupFor(const K& target) const
      {
        return timeouts.find(target) != timeouts.end();
      }

      bool
      HasPendingLookupFrom(const TXOwner& owner) const
      {
        return GetPendingLookupFrom(owner) != nullptr;
      }

      void
      NewTX(const TXOwner& askpeer, const TXOwner& whoasked, const K& k,
            TX< K, V >* t, llarp_time_t requestTimeoutMS = 15s);

      /// mark tx as not fond
      void
      NotFound(const TXOwner& from, const std::unique_ptr< Key_t >& next);

      void
      Found(const TXOwner& from, const K& k, const std::vector< V >& values)
      {
        Inform(from, k, values, true);
      }

      /// inform all watches for key of values found
      void
      Inform(TXOwner from, K key, std::vector< V > values,
             bool sendreply = false, bool removeTimeouts = true);

      void
      Expire(llarp_time_t now);
    };

    template < typename K, typename V, typename K_Hash >
    const TX< K, V >*
    TXHolder< K, V, K_Hash >::GetPendingLookupFrom(const TXOwner& owner) const
    {
      auto itr = tx.find(owner);
      if(itr == tx.end())
      {
        return nullptr;
      }

      return itr->second.get();
    }

    template < typename K, typename V, typename K_Hash >
    void
    TXHolder< K, V, K_Hash >::NewTX(const TXOwner& askpeer,
                                    const TXOwner& whoasked, const K& k,
                                    TX< K, V >* t,
                                    llarp_time_t requestTimeoutMS)
    {
      (void)whoasked;
      tx.emplace(askpeer, std::unique_ptr< TX< K, V > >(t));
      auto count = waiting.count(k);
      waiting.emplace(k, askpeer);

      auto itr = timeouts.find(k);
      if(itr == timeouts.end())
      {
        timeouts.emplace(k, time_now_ms() + requestTimeoutMS);
      }
      if(count == 0)
      {
        t->Start(askpeer);
      }
    }

    template < typename K, typename V, typename K_Hash >
    void
    TXHolder< K, V, K_Hash >::NotFound(const TXOwner& from,
                                       const std::unique_ptr< Key_t >&)
    {
      auto txitr = tx.find(from);
      if(txitr == tx.end())
      {
        return;
      }
      Inform(from, txitr->second->target, {}, true, true);
    }

    template < typename K, typename V, typename K_Hash >
    void
    TXHolder< K, V, K_Hash >::Inform(TXOwner from, K key,
                                     std::vector< V > values, bool sendreply,
                                     bool removeTimeouts)
    {
      auto range = waiting.equal_range(key);
      auto itr   = range.first;
      while(itr != range.second)
      {
        auto txitr = tx.find(itr->second);
        if(txitr != tx.end())
        {
          for(const auto& value : values)
          {
            txitr->second->OnFound(from.node, value);
          }
          if(sendreply)
          {
            txitr->second->SendReply();
            tx.erase(txitr);
          }
        }
        ++itr;
      }

      if(sendreply)
      {
        waiting.erase(key);
      }

      if(removeTimeouts)
      {
        timeouts.erase(key);
      }
    }

    template < typename K, typename V, typename K_Hash >
    void
    TXHolder< K, V, K_Hash >::Expire(llarp_time_t now)
    {
      auto itr = timeouts.begin();
      while(itr != timeouts.end())
      {
        if(now >= itr->second)
        {
          Inform(TXOwner{}, itr->first, {}, true, false);
          itr = timeouts.erase(itr);
        }
        else
        {
          ++itr;
        }
      }
    }
  }  // namespace dht
}  // namespace llarp

#endif