From 7109ddc951680c4175842637c09156dca070bd64 Mon Sep 17 00:00:00 2001
From: Stephen Shelton <steve@brewcraft.org>
Date: Tue, 26 May 2020 13:17:51 -0600
Subject: [PATCH] Add PeerDb::modifyPeerStats()

---
 llarp/peerstats/peer_db.cpp     | 12 +++++++++++-
 llarp/peerstats/peer_db.hpp     | 16 ++++++++++++++++
 test/peerstats/test_peer_db.cpp | 24 ++++++++++++++++++++++++
 3 files changed, 51 insertions(+), 1 deletion(-)

diff --git a/llarp/peerstats/peer_db.cpp b/llarp/peerstats/peer_db.cpp
index e81ea361e..77cb817a6 100644
--- a/llarp/peerstats/peer_db.cpp
+++ b/llarp/peerstats/peer_db.cpp
@@ -100,6 +100,16 @@ namespace llarp
       itr->second += delta;
   }
 
+  void
+  PeerDb::modifyPeerStats(const RouterID& routerId, std::function<void(PeerStats&)> callback)
+  {
+    std::lock_guard gaurd(m_statsLock);
+
+    PeerStats& stats = m_peerStats[routerId];
+    stats.routerId = routerId.ToString();
+    callback(stats);
+  }
+
   std::optional<PeerStats>
   PeerDb::getCurrentPeerStats(const RouterID& routerId) const
   {
@@ -125,7 +135,7 @@ namespace llarp
   bool
   PeerDb::shouldFlush(llarp_time_t now)
   {
-    static constexpr llarp_time_t TargetFlushInterval = 30s;
+    constexpr llarp_time_t TargetFlushInterval = 30s;
     return (now - m_lastFlush.load() >= TargetFlushInterval);
   }
 
diff --git a/llarp/peerstats/peer_db.hpp b/llarp/peerstats/peer_db.hpp
index fb3a8e5f1..9b668d65b 100644
--- a/llarp/peerstats/peer_db.hpp
+++ b/llarp/peerstats/peer_db.hpp
@@ -1,6 +1,7 @@
 #pragma once
 
 #include <filesystem>
+#include <functional>
 #include <unordered_map>
 
 #include <sqlite_orm/sqlite_orm.h>
@@ -54,9 +55,24 @@ namespace llarp
     /// 3) Call accumulatePeerStats() with the stats
     /// 4) Reset the stats to 0
     /// 5) <Repeat 2-4 periodically>
+    ///
+    /// @param routerId is the id of the router whose stats should be modified.
+    /// @param delta is the stats to add to the existing stats
     void
     accumulatePeerStats(const RouterID& routerId, const PeerStats& delta);
 
+    /// Allows write-access to the stats for a given peer while appropriate mutex lock is held. This
+    /// is an alternative means of incrementing peer stats that is suitable for one-off
+    /// modifications.
+    ///
+    /// Note that this holds m_statsLock during the callback invocation, so the callback should
+    /// return as quickly as possible.
+    ///
+    /// @param routerId is the id of the router whose stats should be modified.
+    /// @param callback is a function which will be called immediately with mutex held
+    void
+    modifyPeerStats(const RouterID& routerId, std::function<void(PeerStats&)> callback);
+
     /// Provides a snapshot of the most recent PeerStats we have for the given peer. If we don't
     /// have any stats for the peer, std::nullopt
     ///
diff --git a/test/peerstats/test_peer_db.cpp b/test/peerstats/test_peer_db.cpp
index 1a07f6b3c..18bc9ecd5 100644
--- a/test/peerstats/test_peer_db.cpp
+++ b/test/peerstats/test_peer_db.cpp
@@ -87,3 +87,27 @@ TEST_CASE("Test PeerDb file-backed database reloads properly", "[PeerDb]")
 
   fs::remove(filename);
 }
+
+TEST_CASE("Test PeerDb modifyPeerStats", "[PeerDb]")
+{
+  const llarp::RouterID id = llarp::test::makeBuf<llarp::RouterID>(0xF2);
+
+  int numTimesCalled = 0;
+
+  llarp::PeerDb db;
+  db.loadDatabase(std::nullopt);
+
+  db.modifyPeerStats(id, [&](llarp::PeerStats& stats) {
+    numTimesCalled++;
+
+    stats.numPathBuilds += 42;
+  });
+
+  db.flushDatabase();
+
+  CHECK(numTimesCalled == 1);
+
+  auto stats = db.getCurrentPeerStats(id);
+  CHECK(stats.has_value());
+  CHECK(stats.value().numPathBuilds == 42);
+}