main.py 5.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204
  1. from hashlib import sha256
  2. from datetime import datetime
  3. import asyncio
  4. EventId = str
  5. EventIds = list[EventId]
  6. class Event:
  7. def __init__(self, parents: EventIds):
  8. self.timestamp = datetime.now().timestamp
  9. self.parents = parents
  10. def set_timestamp(self, timestamp):
  11. self.timestamp = timestamp
  12. def hash(self) -> str:
  13. m = sha256()
  14. m.update(str.encode(str(self.timestamp)))
  15. for p in self.parents:
  16. m.update(str.encode(str(p)))
  17. return m.digest().hex()
  18. def __str__(self):
  19. res = f"{self.hash()}"
  20. for p in self.parents:
  21. res += f"\n |"
  22. res += f"\n - {p}"
  23. res += f"\n"
  24. return res
  25. """
  26. ## Graph Example
  27. E1: []
  28. E2: [E1]
  29. E3: [E1]
  30. E4: [E3]
  31. E5: [E3]
  32. E6: [E4, E5]
  33. E7: [E4]
  34. E8: [E2]
  35. """
  36. class Graph:
  37. def __init__(self):
  38. self.events = dict()
  39. # NOTE: we will need to keep track of heads for creating new events.
  40. # Not needed for this demo though.
  41. def add_event(self, event: Event):
  42. self.events[event.hash()] = event
  43. def remove_event(self, event_id: EventId):
  44. if event_id in self.events:
  45. del self.events[event_id]
  46. # check if given events are exist in the graph
  47. # return a list of missing events
  48. def check(self, events: EventIds) -> EventIds:
  49. missing_events = []
  50. for e in events:
  51. if e not in self.events:
  52. missing_events.append(e)
  53. return missing_events
  54. def __str__(self):
  55. res = ""
  56. for event in self.events.values():
  57. res += f"\n {event}"
  58. return res
  59. class Node:
  60. def __init__(self, name: str):
  61. self.name = name
  62. self.orphan_pool = Graph()
  63. self.active_pool = Graph()
  64. # the active pool should always start with one event
  65. genesis_event = Event([])
  66. genesis_event.set_timestamp(0.0)
  67. self.genesis_event = genesis_event
  68. self.active_pool.add_event(genesis_event)
  69. def receive_new_event(self, event: Event):
  70. # reject events with no parents
  71. if len(event.parents) == 0:
  72. return
  73. missing_parents = self.active_pool.check(event.parents)
  74. if len(missing_parents) == 0:
  75. # if there are no missing parents
  76. # add the event to active pool
  77. self.active_pool.add_event(event)
  78. # events list to be removed from orphan pool
  79. # after relink
  80. remove_list: EventIds = []
  81. self.relink(event, remove_list)
  82. # clean up orphan pool
  83. for ev in remove_list:
  84. self.orphan_pool.remove_event(ev)
  85. else:
  86. # add the received event to the orphan pool
  87. self.orphan_pool.add_event(event)
  88. # check if all the missing parents are in orphan pool
  89. # if the missing parents and their links not in orphan pool, request
  90. # them from the network
  91. request_list = []
  92. self.check_parents(request_list, missing_parents)
  93. # XXX
  94. # send all the missing parents in request_list
  95. # to the node who send this event
  96. def check_parents(self, request_list, parents: EventIds):
  97. for parent_hash in parents:
  98. if parent_hash in self.orphan_pool.events:
  99. parent = self.orphan_pool.events[parent_hash]
  100. # recursive call
  101. self.check_parents(request_list, parent.parents)
  102. else:
  103. request_list.append(parent_hash)
  104. def relink(self, event: Event, remove_list=[]):
  105. # check if the orphan pool has an event linked
  106. # to the new added event
  107. for (orphan_hash, orphan) in self.orphan_pool.events.items():
  108. if orphan_hash in remove_list:
  109. continue
  110. if event.hash() not in orphan.parents:
  111. continue
  112. missing_parents = self.active_pool.check(orphan.parents)
  113. if len(missing_parents) == 0:
  114. self.active_pool.add_event(orphan)
  115. remove_list.append(orphan_hash)
  116. # recursive call
  117. self.relink(orphan, remove_list)
  118. def __str__(self):
  119. return f"""------
  120. \n Name: {self.name}
  121. \n Active Pool: {self.active_pool}
  122. \n Orphan Pool: {self.orphan_pool}"""
  123. async def run_node(name):
  124. print(f"{name} Started")
  125. node = Node(name)
  126. print(f"{name} End")
  127. async def main():
  128. tasks = await asyncio.gather(
  129. run_node("NodeA"),
  130. run_node("NodeB"),
  131. run_node("NodeC"))
  132. def test_node():
  133. node_a = Node("NodeA")
  134. event0 = node_a.genesis_event
  135. event1 = Event([event0.hash()])
  136. event2 = Event([event1.hash()])
  137. event3 = Event([event2.hash(), event0.hash()])
  138. event4 = Event([event1.hash(), event3.hash()])
  139. event5 = Event([event4.hash(), "FAKEHASH"])
  140. event6 = Event([event5.hash(), event3.hash()])
  141. node_a.receive_new_event(event3)
  142. node_a.receive_new_event(event2)
  143. node_a.receive_new_event(event1)
  144. node_a.receive_new_event(event5)
  145. node_a.receive_new_event(event6)
  146. node_a.receive_new_event(event4)
  147. print(node_a)
  148. if __name__ == "__main__":
  149. test_node()
  150. asyncio.run(main())