Short-circuit balance cache

apps/contract-migration/scripts/traffic.py

@@ -225,29 +225,41 @@ class TrafficTasker:
         self.aux[k] = v
 
 
-    def balances(self, accounts=None):
+    def balances(self, accounts=None, refresh=False):
-        if accounts == None:
+        if refresh:
-            accounts = self.accounts
+            if accounts == None:
-        for account in accounts:
+                accounts = self.accounts
-            for token in self.tokens:
+            for account in accounts:
-                # TODO: use proper redis callback 
+                for token in self.tokens:
-                api = Api(
+                    value = self.balance(account, token)
-                    str(self.aux['chain_spec']),
+                    self.__cache_balance(account, token.symbol(), value)
-                    queue=self.aux['api_queue'],
+                    logg.debug('balance sender {} token {} = {}'.format(account, token, value))
-                    #callback_param='{}:{}:{}:{}'.format(aux['redis_host_callback'], aux['redis_port_callback'], aux['redis_db'], aux['redis_channel']),
+        else:
-                    #callback_task='cic_eth.callbacks.redis.redis',
+            logg.debug('returning cached balances')
-                    #callback_queue=queue,
-                    )
-                t = api.balance(account, token.symbol())
-                r = t.get()
-                for c in t.collect():
-                    r = c[1]
-                assert t.successful()
-                self.__cache_balance(account, token.symbol(), r[0])
-                logg.debug('balance sender {} token {} = {}'.format(account, token, r))
         return self.__balances
 
 
+    def balance(self, account, token):
+        # TODO: use proper redis callback 
+        api = Api(
+            str(self.aux['chain_spec']),
+            queue=self.aux['api_queue'],
+            #callback_param='{}:{}:{}:{}'.format(aux['redis_host_callback'], aux['redis_port_callback'], aux['redis_db'], aux['redis_channel']),
+            #callback_task='cic_eth.callbacks.redis.redis',
+            #callback_queue=queue,
+            )
+        t = api.balance(account, token.symbol())
+        r = t.get()
+        for c in t.collect():
+            r = c[1]
+        assert t.successful()
+        return r[0]
+
+
+    def update_balance(self, account, token, value):
+        self.__cache_balance(account, token.symbol(), value)
+
+
 class Handler:
 
     def __init__(self, config, traffic_router):
@@ -256,6 +268,7 @@ class Handler:
         self.pubsub = self.__connect_redis(self.redis_channel, config)
         self.traffic_items = {}
         self.config = config
+        self.init = False
 
 
     def __connect_redis(self, redis_channel, config):
@@ -271,6 +284,10 @@ class Handler:
         traffic_tasker = TrafficTasker()
         traffic_tasker.add_aux('redis_channel', self.redis_channel)
 
+        refresh_balance = not self.init
+        balances = traffic_tasker.balances(refresh=refresh_balance)
+        self.init = True
+
         if len(traffic_tasker.tokens) == 0:
             logg.error('patiently waiting for at least one registered token...')
             return
@@ -278,10 +295,8 @@ class Handler:
         logg.debug('executing handler refresh with accouts {}'.format(traffic_tasker.accounts))
         logg.debug('executing handler refresh with tokens {}'.format(traffic_tasker.tokens))
 
-        #senders = copy.copy(traffic_tasker.accounts)
         sender_indices = [*range(0, len(traffic_tasker.accounts))]
         # TODO: only get balances for the selection that we will be generating for
-        balances = traffic_tasker.balances()
 
         while True:
             traffic_item = traffic_router.reserve()
@@ -297,17 +312,18 @@ class Handler:
             sender_index_index = random.randint(0, len(sender_indices)-1)
             sender_index = sender_indices[sender_index_index]
             sender = traffic_tasker.accounts[sender_index]
-            logg.debug('balances {}'.format(balances))
+            #balance_full = balances[sender][token_pair[0].symbol()]
-            balance_unit = balances[sender][token_pair[0].symbol()]
-            balance = balance_unit['balance_network'] - balance_unit['balance_outgoing']
             if len(sender_indices) == 1:
                 sender_indices[m] = sender_sender_indices[len(senders)-1]
             sender_indices = sender_indices[:len(sender_indices)-1]
 
-            recipient_index = random.randint(0, len(traffic_tasker.accounts))
+            balance_full = traffic_tasker.balance(sender, token_pair[0])
+            balance = balance_full['balance_network'] - balance_full['balance_outgoing']
+
+            recipient_index = random.randint(0, len(traffic_tasker.accounts)-1)
             recipient = traffic_tasker.accounts[recipient_index]
             
-            (e, t, balance,) = traffic_item.method(
+            (e, t, balance_result,) = traffic_item.method(
                     token_pair,
                     sender,
                     recipient,
@@ -316,7 +332,7 @@ class Handler:
                     block_number,
                     tx_index,
                     )
-            balances[sender][token_pair[0].symbol()] = balance_unit
+            traffic_tasker.update_balance(sender, token_pair[0], balance_result)
             sender_indices.append(recipient_index)
 
             if e != None:
@@ -331,7 +347,6 @@ class Handler:
             self.traffic_items[traffic_item.ext] = traffic_item
 
 
-        # TODO: add drain
         while True:
             m = self.pubsub.get_message(timeout=0.1)
             if m == None: