Add E3RequestComplete event and kill actors on completion

packages/ciphernode/aggregator/src/plaintext_aggregator.rs

@@ -1,7 +1,7 @@
 use actix::prelude::*;
 use anyhow::Result;
 use enclave_core::{
-    DecryptionshareCreated, E3id, EnclaveEvent, EventBus, OrderedSet, PlaintextAggregated, Seed,
+    DecryptionshareCreated, Die, E3id, EnclaveEvent, EventBus, OrderedSet, PlaintextAggregated, E3RequestComplete, Seed
 };
 use fhe::{Fhe, GetAggregatePlaintext};
 use sortition::{GetHasNode, Sortition};
@@ -191,6 +191,14 @@ impl Handler<ComputeAggregate> for PlaintextAggregator {
 
         self.bus.do_send(event);
 
+
         Ok(())
     }
 }
+
+impl Handler<Die> for PlaintextAggregator {
+    type Result = ();
+    fn handle(&mut self, _: Die, ctx: &mut Self::Context) -> Self::Result {
+       ctx.stop()
+    }
+}

packages/ciphernode/aggregator/src/publickey_aggregator.rs

@@ -1,7 +1,7 @@
 use actix::prelude::*;
 use anyhow::Result;
 use enclave_core::{
-    E3id, EnclaveEvent, EventBus, KeyshareCreated, OrderedSet, PublicKeyAggregated, Seed,
+    Die, E3id, EnclaveEvent, EventBus, KeyshareCreated, OrderedSet, PublicKeyAggregated, Seed
 };
 use fhe::{Fhe, GetAggregatePublicKey};
 use sortition::{GetHasNode, GetNodes, Sortition};
@@ -219,3 +219,10 @@ impl Handler<NotifyNetwork> for PublicKeyAggregator {
         )
     }
 }
+
+impl Handler<Die> for PublicKeyAggregator {
+    type Result = ();
+    fn handle(&mut self, _: Die, ctx: &mut Self::Context) -> Self::Result {
+       ctx.stop()
+    }
+}

packages/ciphernode/core/src/events.rs

@@ -1,4 +1,4 @@
-use actix::Message;
+use actix::{Actor, Addr, Message};
 use alloy::{
     hex,
     primitives::{Uint, U256},
@@ -116,6 +116,10 @@ pub enum EnclaveEvent {
         id: EventId,
         data: EnclaveError,
     },
+    E3RequestComplete {
+        id: EventId,
+        data: E3RequestComplete,
+    },
     // CommitteeSelected,
     // OutputDecrypted,
     // CiphernodeRegistered,
@@ -142,6 +146,7 @@ impl EnclaveEvent {
             EnclaveEvent::E3Requested { .. } => true,
             EnclaveEvent::CiphernodeAdded { .. } => true,
             EnclaveEvent::CiphernodeRemoved { .. } => true,
+            EnclaveEvent::E3RequestComplete { .. } => true,
             _ => false,
         }
     }
@@ -160,6 +165,7 @@ impl From<EnclaveEvent> for EventId {
             EnclaveEvent::CiphernodeAdded { id, .. } => id,
             EnclaveEvent::CiphernodeRemoved { id, .. } => id,
             EnclaveEvent::EnclaveError { id, .. } => id,
+            EnclaveEvent::E3RequestComplete { id, .. } => id,
         }
     }
 }
@@ -238,6 +244,15 @@ impl From<PlaintextAggregated> for EnclaveEvent {
     }
 }
 
+impl From<E3RequestComplete> for EnclaveEvent {
+    fn from(data: E3RequestComplete) -> Self {
+        EnclaveEvent::E3RequestComplete {
+            id: EventId::from(data.clone()),
+            data: data.clone(),
+        }
+    }
+}
+
 impl From<CiphernodeSelected> for EnclaveEvent {
     fn from(data: CiphernodeSelected) -> Self {
         EnclaveEvent::CiphernodeSelected {
@@ -349,6 +364,14 @@ pub struct PlaintextAggregated {
     pub src_chain_id: u64,
 }
 
+
+/// E3RequestComplete event is a local only event
+#[derive(Message, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[rtype(result = "()")]
+pub struct E3RequestComplete {
+    pub e3_id: E3id,
+}
+
 #[derive(Message, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
 #[rtype(result = "()")]
 pub struct CiphernodeAdded {
@@ -372,6 +395,10 @@ pub struct EnclaveError {
     pub message: String,
 }
 
+#[derive(Message, Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
+#[rtype(result = "()")]
+pub struct Die;
+
 #[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, Serialize, Deserialize)]
 pub struct Seed(pub [u8; 32]);
 impl From<Seed> for u64 {

packages/ciphernode/keyshare/src/keyshare.rs

@@ -2,8 +2,7 @@ use actix::prelude::*;
 use anyhow::{anyhow, Context, Result};
 use data::{Data, Get, Insert};
 use enclave_core::{
-    CiphernodeSelected, CiphertextOutputPublished, DecryptionshareCreated, EnclaveErrorType,
-    EnclaveEvent, EventBus, FromError, KeyshareCreated,
+    CiphernodeSelected, CiphertextOutputPublished, DecryptionshareCreated, Die, EnclaveErrorType, EnclaveEvent, EventBus, FromError, KeyshareCreated
 };
 use fhe::{DecryptCiphertext, Fhe};
 use std::sync::Arc;
@@ -99,6 +98,13 @@ impl Handler<CiphertextOutputPublished> for Keyshare {
     }
 }
 
+impl Handler<Die> for Keyshare {
+    type Result = ();
+    fn handle(&mut self, _: Die, ctx: &mut Self::Context) -> Self::Result {
+       ctx.stop()
+    }
+}
+
 async fn on_decryption_requested(
     fhe: Arc<Fhe>,
     data: Addr<Data>,

packages/ciphernode/router/src/e3_request_router.rs

@@ -3,6 +3,7 @@ use crate::CommitteeMetaFactory;
 use super::CommitteeMeta;
 use aggregator::PlaintextAggregator;
 use aggregator::PublicKeyAggregator;
+use enclave_core::E3RequestComplete;
 use enclave_core::{E3id, EnclaveEvent, EventBus, Subscribe};
 use fhe::Fhe;
 use keyshare::Keyshare;
@@ -85,6 +86,7 @@ pub struct E3RequestRouter {
     contexts: HashMap<E3id, E3RequestContext>,
     hooks: Vec<EventHook>,
     buffer: EventBuffer,
+    bus: Addr<EventBus>,
 }
 
 impl E3RequestRouter {
@@ -107,13 +109,26 @@ impl Handler<EnclaveEvent> for E3RequestRouter {
             return;
         };
 
-        let context = self.contexts.entry(e3_id).or_default();
+        let context = self.contexts.entry(e3_id.clone()).or_default();
 
         for hook in &mut self.hooks {
             hook(context, msg.clone());
         }
 
         context.forward_message(&msg, &mut self.buffer);
+
+        // Here we are detemining that by receiving the PlaintextAggregated event our request is
+        // complete and we can notify everyone. This might change as we consider other factors
+        // when determining if the request is complete
+        if let EnclaveEvent::PlaintextAggregated { .. } = msg {
+            // Local event to clean up context
+            let event = EnclaveEvent::from(E3RequestComplete {
+                e3_id: e3_id.clone(),
+            });
+
+            // Send to bus so all other actors can react to a request being complete.
+            self.bus.do_send(event);
+        }
     }
 }
 
@@ -133,6 +148,7 @@ impl E3RequestRouterBuilder {
             contexts: HashMap::new(),
             hooks: self.hooks,
             buffer: EventBuffer::default(),
+            bus: self.bus.clone()
         };
 
         let addr = e3r.start();

packages/ciphernode/router/src/hooks.rs

@@ -2,7 +2,7 @@ use crate::EventHook;
 use actix::{Actor, Addr};
 use aggregator::{PlaintextAggregator, PublicKeyAggregator};
 use data::Data;
-use enclave_core::{E3Requested, EnclaveEvent, EventBus};
+use enclave_core::{Die, E3Requested, EnclaveEvent, EventBus};
 use fhe::{Fhe, SharedRng};
 use keyshare::Keyshare;
 use sortition::Sortition;
@@ -30,84 +30,103 @@ pub struct LazyKeyshare;
 impl LazyKeyshare {
     pub fn create(bus: Addr<EventBus>, data: Addr<Data>, address: &str) -> EventHook {
         let address = address.to_string();
-        Box::new(move |ctx, evt| {
-            // Save Ciphernode on CiphernodeSelected
-            let EnclaveEvent::CiphernodeSelected { .. } = evt else {
-                return;
-            };
+        Box::new(move |ctx, evt| match evt {
+            EnclaveEvent::CiphernodeSelected { .. } => {
+                let Some(ref fhe) = ctx.fhe else {
+                    return;
+                };
 
-            let Some(ref fhe) = ctx.fhe else {
-                return;
-            };
+                ctx.keyshare =
+                    Some(Keyshare::new(bus.clone(), data.clone(), fhe.clone(), &address).start())
+            }
+            EnclaveEvent::E3RequestComplete { .. } => {
+                let Some(actor) = ctx.keyshare.take() else {
+                    return;
+                };
 
-            ctx.keyshare =
-                Some(Keyshare::new(bus.clone(), data.clone(), fhe.clone(), &address).start())
+                actor.do_send(Die);
+            }
+            _ => (),
         })
     }
 }
 
 pub struct LazyPlaintextAggregator;
 impl LazyPlaintextAggregator {
     pub fn create(bus: Addr<EventBus>, sortition: Addr<Sortition>) -> EventHook {
-        Box::new(move |ctx, evt| {
-            // Save plaintext aggregator
-            let EnclaveEvent::CiphertextOutputPublished { data, .. } = evt else {
-                return;
-            };
-            let Some(ref fhe) = ctx.fhe else {
-                return;
-            };
-            let Some(ref meta) = ctx.meta else {
-                return;
-            };
+        Box::new(move |ctx, evt| match evt {
+            EnclaveEvent::CiphertextOutputPublished { data, .. } => {
+                let Some(ref fhe) = ctx.fhe else {
+                    return;
+                };
+                let Some(ref meta) = ctx.meta else {
+                    return;
+                };
+
+                // Save plaintext aggregator
+                ctx.plaintext = Some(
+                    PlaintextAggregator::new(
+                        fhe.clone(),
+                        bus.clone(),
+                        sortition.clone(),
+                        data.e3_id,
+                        meta.threshold_m,
+                        meta.seed,
+                        data.ciphertext_output,
+                        meta.src_chain_id,
+                    )
+                    .start(),
+                );
+            }
+            EnclaveEvent::E3RequestComplete { .. } => {
+                let Some(actor) = ctx.plaintext.take() else {
+                    return;
+                };
 
-            ctx.plaintext = Some(
-                PlaintextAggregator::new(
-                    fhe.clone(),
-                    bus.clone(),
-                    sortition.clone(),
-                    data.e3_id,
-                    meta.threshold_m,
-                    meta.seed,
-                    data.ciphertext_output,
-                    meta.src_chain_id,
-                )
-                .start(),
-            );
+                actor.do_send(Die);
+            }
+            _ => (),
         })
     }
 }
 
 pub struct LazyPublicKeyAggregator;
 impl LazyPublicKeyAggregator {
     pub fn create(bus: Addr<EventBus>, sortition: Addr<Sortition>) -> EventHook {
-        Box::new(move |ctx, evt| {
+        Box::new(move |ctx, evt| match evt {
             // Saving the publickey aggregator with deps on E3Requested
-            let EnclaveEvent::E3Requested { data, .. } = evt else {
-                return;
-            };
+            EnclaveEvent::E3Requested { data, .. } => {
+                let Some(ref fhe) = ctx.fhe else {
+                    println!("fhe was not on ctx");
+                    return;
+                };
+                let Some(ref meta) = ctx.meta else {
+                    println!("meta was not on ctx");
+                    return;
+                };
 
-            let Some(ref fhe) = ctx.fhe else {
-                println!("fhe was not on ctx");
-                return;
-            };
-            let Some(ref meta) = ctx.meta else {
-                println!("meta was not on ctx");
-                return;
-            };
+                ctx.publickey = Some(
+                    PublicKeyAggregator::new(
+                        fhe.clone(),
+                        bus.clone(),
+                        sortition.clone(),
+                        data.e3_id,
+                        meta.threshold_m,
+                        meta.seed,
+                        meta.src_chain_id,
+                    )
+                    .start(),
+                );
+            }
+            EnclaveEvent::E3RequestComplete { .. } => {
+                let Some(actor) = ctx.publickey.take() else {
+                    return;
+                };
+
+                actor.do_send(Die);
+            }
 
-            ctx.publickey = Some(
-                PublicKeyAggregator::new(
-                    fhe.clone(),
-                    bus.clone(),
-                    sortition.clone(),
-                    data.e3_id,
-                    meta.threshold_m,
-                    meta.seed,
-                    meta.src_chain_id,
-                )
-                .start(),
-            );
+            _ => (),
         })
     }
 }