From 78f0d93a270a6731d305025a806aa70e1810fee7 Mon Sep 17 00:00:00 2001
From: derekpierre <derek.pierre@gmail.com>
Date: Fri, 9 Aug 2024 08:18:08 -0400
Subject: [PATCH] Fix bug where node doesn't have a bucket when we try to mark
 it as unsuccessful. Annotate what node dictionaries represent.

---
 porter/main.py | 13 ++++++++-----
 1 file changed, 8 insertions(+), 5 deletions(-)

diff --git a/porter/main.py b/porter/main.py
index 22a104f..767a69a 100644
--- a/porter/main.py
+++ b/porter/main.py
@@ -392,8 +392,10 @@ def __init__(self, _reservoir, need_successes: int):
                 self.reservoir = _reservoir
                 self.need_successes = need_successes
                 self.predefined_buckets = self.read_buckets()
-                self.bucketed_nodes = defaultdict(list)
-                self.selected_nodes = dict()
+                self.bucketed_nodes = defaultdict(
+                    list
+                )  # <bucket> -> <list of checksum addresses>
+                self.selected_nodes = dict()  # <checksum address> -> <bucket>
 
             def read_buckets(self) -> Dict:
                 try:
@@ -420,9 +422,10 @@ def find_bucket(self, node):
                         return bucket_name
                 return None
 
-            def mark_as_not_successful(self, failure: ChecksumAddress):
-                bucket = self.selected_nodes[failure]
-                self.bucketed_nodes[bucket].remove(failure)
+            def mark_as_not_successful(self, unsuccessful_node: ChecksumAddress):
+                bucket = self.selected_nodes.get(unsuccessful_node)
+                if bucket:
+                    self.bucketed_nodes[bucket].remove(unsuccessful_node)
 
             def __call__(self, _successes: int) -> Optional[List[ChecksumAddress]]:
                 batch = []