From d3af95d1eba1725b21370447f6d02611c010ecf6 Mon Sep 17 00:00:00 2001
From: Age Manning <Age@AgeManning.com>
Date: Wed, 27 Mar 2019 22:41:55 +1100
Subject: [PATCH] Returns attestation duty for validator client processing

---
 validator_client/src/duties/epoch_duties.rs | 34 +++++++++++++--------
 validator_client/src/duties/mod.rs          |  8 ++---
 validator_client/src/service.rs             |  5 ++-
 3 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/validator_client/src/duties/epoch_duties.rs b/validator_client/src/duties/epoch_duties.rs
index a74ded18b..5c23e82b1 100644
--- a/validator_client/src/duties/epoch_duties.rs
+++ b/validator_client/src/duties/epoch_duties.rs
@@ -1,12 +1,17 @@
 use std::collections::HashMap;
 use std::ops::{Deref, DerefMut};
-use types::{Epoch, PublicKey, Slot};
+use types::{AttestationDuty, Epoch, PublicKey, Slot};
 
-/// The type of work a validator is required to do in a given slot.
+/// When work needs to be performed by a validator, this type is given back to the main service
+/// which indicates all the information that required to process the work.
+///
+/// Note: This is calculated per slot, so a validator knows which slot is related to this struct.
 #[derive(Debug, Clone)]
-pub struct WorkType {
+pub struct WorkInfo {
+    /// Validator needs to produce a block.
     pub produce_block: bool,
-    pub produce_attestation: bool,
+    /// Validator needs to produce an attestation. This supplies the required attestation data.
+    pub attestation_duty: Option<AttestationDuty>,
 }
 
 /// The information required for a validator to propose and attest during some epoch.
@@ -23,23 +28,28 @@ pub struct EpochDuty {
 }
 
 impl EpochDuty {
-    /// Returns `WorkType` if work needs to be done in the supplied `slot`
-    pub fn is_work_slot(&self, slot: Slot) -> Option<WorkType> {
+    /// Returns `WorkInfo` if work needs to be done in the supplied `slot`
+    pub fn is_work_slot(&self, slot: Slot) -> Option<WorkInfo> {
         // if validator is required to produce a slot return true
         let produce_block = match self.block_production_slot {
             Some(s) if s == slot => true,
             _ => false,
         };
 
-        let mut produce_attestation = false;
+        // if the validator is required to attest to a shard, create the data
+        let mut attestation_duty = None;
         if self.attestation_slot == slot {
-            produce_attestation = true;
+            attestation_duty = Some(AttestationDuty {
+                slot,
+                shard: self.attestation_shard,
+                committee_index: self.committee_index as usize,
+            });
         }
 
-        if produce_block | produce_attestation {
-            return Some(WorkType {
+        if produce_block | attestation_duty.is_some() {
+            return Some(WorkInfo {
                 produce_block,
-                produce_attestation,
+                attestation_duty,
             });
         }
         None
@@ -89,7 +99,7 @@ impl EpochDutiesMap {
         &self,
         slot: Slot,
         pubkey: &PublicKey,
-    ) -> Result<Option<WorkType>, EpochDutiesMapError> {
+    ) -> Result<Option<WorkInfo>, EpochDutiesMapError> {
         let epoch = slot.epoch(self.slots_per_epoch);
 
         let epoch_duties = self
diff --git a/validator_client/src/duties/mod.rs b/validator_client/src/duties/mod.rs
index 94b845083..51470827c 100644
--- a/validator_client/src/duties/mod.rs
+++ b/validator_client/src/duties/mod.rs
@@ -6,7 +6,7 @@ mod grpc;
 mod traits;
 
 use self::epoch_duties::{EpochDuties, EpochDutiesMapError};
-pub use self::epoch_duties::{EpochDutiesMap, WorkType};
+pub use self::epoch_duties::{EpochDutiesMap, WorkInfo};
 use self::traits::{BeaconNode, BeaconNodeError};
 use futures::Async;
 use slog::{debug, error, info};
@@ -85,10 +85,10 @@ impl<U: BeaconNode> DutiesManager<U> {
         Ok(Async::Ready(()))
     }
 
-    /// Returns a list of (Public, WorkType) indicating all the validators that have work to perform
+    /// Returns a list of (Public, WorkInfo) indicating all the validators that have work to perform
     /// this slot.
-    pub fn get_current_work(&self, slot: Slot) -> Option<Vec<(PublicKey, WorkType)>> {
-        let mut current_work: Vec<(PublicKey, WorkType)> = Vec::new();
+    pub fn get_current_work(&self, slot: Slot) -> Option<Vec<(PublicKey, WorkInfo)>> {
+        let mut current_work: Vec<(PublicKey, WorkInfo)> = Vec::new();
 
         // if the map is poisoned, return None
         let duties = self.duties_map.read().ok()?;
diff --git a/validator_client/src/service.rs b/validator_client/src/service.rs
index 07d88d344..f62ade97e 100644
--- a/validator_client/src/service.rs
+++ b/validator_client/src/service.rs
@@ -272,7 +272,10 @@ impl Service {
                         if work_type.produce_block {
                             // TODO: Produce a beacon block in a new thread
                         }
-                        if work_type.produce_attestation {
+                        if work_type.attestation_duty.is_some() {
+                            // available AttestationDuty info
+                            let attestation_duty =
+                                work_type.attestation_duty.expect("Cannot be None");
                             //TODO: Produce an attestation in a new thread
                         }
                     }