From f31cf9a656932e272a639088bbbae265fcb62253 Mon Sep 17 00:00:00 2001
From: ptrus <ptrusr@gmail.com>
Date: Mon, 20 Apr 2020 12:38:50 +0200
Subject: [PATCH] go/extra/stats: fix & simplify node-entity mapping

Instead of separately querying for entities and nodes, we can get Entity IDs
from nodes directly.

This change also fixes a case that previous variant missed: node that was
removed from entity list of nodes, but has not yet expired.
---
 .changelog/2856.bugfix.md   |  7 +++++
 go/extra/stats/cmd/stats.go | 55 ++++++++++---------------------------
 2 files changed, 21 insertions(+), 41 deletions(-)
 create mode 100644 .changelog/2856.bugfix.md

diff --git a/.changelog/2856.bugfix.md b/.changelog/2856.bugfix.md
new file mode 100644
index 00000000000..e21d58918d0
--- /dev/null
+++ b/.changelog/2856.bugfix.md
@@ -0,0 +1,7 @@
+go/extra/stats: fix & simplify node-entity mapping
+
+Instead of separately querying for entities and nodes, we can get Entity IDs
+from nodes directly.
+
+This change also fixes a case that previous variant missed: node that was
+removed from entity list of nodes, but has not yet expired.
diff --git a/go/extra/stats/cmd/stats.go b/go/extra/stats/cmd/stats.go
index e904e613c46..a25faa4dbb1 100644
--- a/go/extra/stats/cmd/stats.go
+++ b/go/extra/stats/cmd/stats.go
@@ -14,7 +14,6 @@ import (
 	"github.com/oasislabs/oasis-core/go/common/cbor"
 	"github.com/oasislabs/oasis-core/go/common/crypto/signature"
 	"github.com/oasislabs/oasis-core/go/common/logging"
-	"github.com/oasislabs/oasis-core/go/common/node"
 	consensusAPI "github.com/oasislabs/oasis-core/go/consensus/api"
 	tmApi "github.com/oasislabs/oasis-core/go/consensus/tendermint/api"
 	tmcrypto "github.com/oasislabs/oasis-core/go/consensus/tendermint/crypto"
@@ -155,62 +154,36 @@ func newStats() *stats {
 }
 
 func (s *stats) addRegistryData(ctx context.Context, registry registryAPI.Backend, height int64) error {
-	// Fetch entities.
-	entities, err := registry.GetEntities(ctx, height)
-	if err != nil {
-		return err
-	}
-
 	// Fetch nodes.
 	nodes, err := registry.GetNodes(ctx, height)
 	if err != nil {
 		return err
 	}
 
-	// Map: nodeID -> Node
-	nodesMap := make(map[signature.PublicKey]*node.Node)
+	// Update stored mappings.
 	for _, n := range nodes {
-		nodesMap[n.ID] = n
-	}
-
-	// Store new nodes and entities info.
-	for _, ent := range entities {
 		var es *entityStats
 		var ok bool
-
-		// Since registry data can be fetched at multiple heights, entities might
-		// already exist.
-		if es, ok = s.entities[ent.ID]; !ok {
+		// Get or create node entity.
+		if es, ok = s.entities[n.EntityID]; !ok {
 			es = &entityStats{
-				id:             ent.ID,
+				id:             n.EntityID,
 				nodeSignatures: make(map[signature.PublicKey]int64),
 				nodeProposals:  make(map[signature.PublicKey]int64),
 			}
-			s.entities[ent.ID] = es
+			s.entities[n.EntityID] = es
 		}
 
-		for _, nodeID := range ent.Nodes {
-			node, ok := nodesMap[nodeID]
-			if !ok {
-				// Skip entity nodes that are not registered.
-				logger.Debug("skipping not registered entity node",
-					"node_id", nodeID,
-				)
-				continue
-			}
-
-			// Add missing nodes.
-			if _, ok := es.nodeSignatures[nodeID]; !ok {
-				es.nodeSignatures[nodeID] = 0
-				es.nodeProposals[nodeID] = 0
-				cID := node.Consensus.ID
-				tmADdr := tmcrypto.PublicKeyToTendermint(&cID).Address().String()
-				s.nodeAddressMap[tmADdr] = nodeIDs{
-					entityID: ent.ID,
-					nodeID:   nodeID,
-				}
+		// Initialize node stats if missing.
+		if _, ok := es.nodeSignatures[n.ID]; !ok {
+			es.nodeSignatures[n.ID] = 0
+			es.nodeProposals[n.ID] = 0
+			cID := n.Consensus.ID
+			tmAddr := tmcrypto.PublicKeyToTendermint(&cID).Address().String()
+			s.nodeAddressMap[tmAddr] = nodeIDs{
+				entityID: n.EntityID,
+				nodeID:   n.ID,
 			}
-
 		}
 	}