Add mass metric

schlafana.py

@@ -5,7 +5,7 @@ import logging
 from pprint import pprint
 from websocket import create_connection
 from prometheus_client import Counter, Gauge, Summary, start_http_server
-from prometheus_client.core import CounterMetricFamily, REGISTRY
+from prometheus_client.core import GaugeMetricFamily, CounterMetricFamily, REGISTRY
 import time
 
 MESSAGES = { "GameInfo" : 0x00,
@@ -44,6 +44,16 @@ class BotRevisionCountCollector(object):
         yield c
 
 
+class BotMassGaugeCollector(object):
+    def collect(self):
+        c = GaugeMetricFamily("mass", "Mass of each snake", labels=["name", "snake_id", "code_id"])
+        for user in game.users.users:
+            user = game.users.users.get(user)
+            rev = user.revs.get(user.current_rev)
+            c.add_metric([user.name, str(user.current_worm), str(user.current_rev)], rev.mass)
+        yield c
+
+
 def main():
     global game
     game = Game()
@@ -61,6 +71,7 @@ def main():
     REGISTRY.register(BotDeathCollector())
     REGISTRY.register(BotKillCollector())
     REGISTRY.register(BotRevisionCountCollector())
+    REGISTRY.register(BotMassGaugeCollector())
     start_http_server(9000)
     logger.info("Start prometheus scrape endpoint")
 
@@ -177,14 +188,7 @@ class BotUser():
 
     def update(self, data):
         rev_id = data.get("db_id")
-        worm_id = data.get("id") #TODO: does this work?
-        if "data" in data.keys():
-            user = game.users.get_user_by_worm(worm_id)
-            current_rev = user.revs.get(user.current_rev)
-            current_rev.carrion_food_consumed = data.get("data").get("c")
-            current_rev.natural_food_consumed = data.get("data").get("n")
-            current_rev.hunted_food_consumed = data.get("data").get("h")
-            pdb.set_trace()
+        worm_id = data.get("id")
         if rev_id in self.revs.keys(): # revision already known
             rev = self.revs.get(rev_id)
             rev.update(data)
@@ -212,8 +216,12 @@ class BotRevision():
         pass
 
     def update(self, data):
-        if "food" in data.keys():
-            pdb.set_trace()
+        if "data" in data.keys():
+            self.mass = data.get("data").get("m")
+            self.carrion_food_consumed = data.get("data").get("c")
+            self.natural_food_consumed = data.get("data").get("n")
+            self.hunted_food_consumed = data.get("data").get("h")
+        else:
             self.current_worm = data.get("id")
             self.rev_id = data.get("db_id")
             self.mass = data.get("mass")