From b4e733f2fafaedf0ccc7f2215686d90b63285c4b Mon Sep 17 00:00:00 2001
From: Warren He <w@oasislabs.com>
Date: Tue, 7 Jan 2020 12:54:55 -0800
Subject: [PATCH] go scheduler: merge and storage eligibility by runtime

---
 .../tendermint/apps/scheduler/scheduler.go    | 22 +++++++++++++++++--
 go/registry/api/api.go                        | 11 +++++-----
 2 files changed, 26 insertions(+), 7 deletions(-)

diff --git a/go/consensus/tendermint/apps/scheduler/scheduler.go b/go/consensus/tendermint/apps/scheduler/scheduler.go
index 8c749db8f13..f41bbc608b7 100644
--- a/go/consensus/tendermint/apps/scheduler/scheduler.go
+++ b/go/consensus/tendermint/apps/scheduler/scheduler.go
@@ -376,7 +376,16 @@ func (app *schedulerApplication) isSuitableComputeWorker(n *node.Node, rt *regis
 }
 
 func (app *schedulerApplication) isSuitableStorageWorker(n *node.Node, rt *registry.Runtime, ts time.Time) bool {
-	return n.HasRoles(node.RoleStorageWorker)
+	if !n.HasRoles(node.RoleStorageWorker) {
+		return false
+	}
+	for _, nrt := range n.Runtimes {
+		if !nrt.ID.Equal(&rt.ID) {
+			continue
+		}
+		return true
+	}
+	return false
 }
 
 func (app *schedulerApplication) isSuitableTransactionScheduler(n *node.Node, rt *registry.Runtime, ts time.Time) bool {
@@ -393,7 +402,16 @@ func (app *schedulerApplication) isSuitableTransactionScheduler(n *node.Node, rt
 }
 
 func (app *schedulerApplication) isSuitableMergeWorker(n *node.Node, rt *registry.Runtime, ts time.Time) bool {
-	return n.HasRoles(node.RoleComputeWorker)
+	if !n.HasRoles(node.RoleComputeWorker) {
+		return false
+	}
+	for _, nrt := range n.Runtimes {
+		if !nrt.ID.Equal(&rt.ID) {
+			continue
+		}
+		return true
+	}
+	return false
 }
 
 // Operates on consensus connection.
diff --git a/go/registry/api/api.go b/go/registry/api/api.go
index febb5ae9451..cb5c48e0cce 100644
--- a/go/registry/api/api.go
+++ b/go/registry/api/api.go
@@ -143,6 +143,7 @@ var (
 
 	// RuntimesRequiredRoles are the Node roles that require runtimes.
 	RuntimesRequiredRoles = node.RoleComputeWorker |
+		node.RoleStorageWorker |
 		node.RoleKeyManager
 
 	// ConsensusAddressRequiredRoles are the Node roles that require Consensus Address.
@@ -1145,15 +1146,15 @@ func SanityCheckNodes(nodes []*node.SignedNode, seenEntities map[signature.Publi
 			return fmt.Errorf("registry: sanity check failed: compute worker node must have runtime(s)")
 		}
 
-		if n.HasRoles(node.RoleKeyManager) && len(n.Runtimes) == 0 {
-			return fmt.Errorf("registry: sanity check failed: key manager node must have runtime(s)")
+		if n.HasRoles(node.RoleStorageWorker) && len(n.Runtimes) == 0 {
+			return fmt.Errorf("registry: sanity check failed: storage worker node must have runtime(s)")
 		}
 
-		if n.HasRoles(node.RoleStorageWorker) && !n.HasRoles(node.RoleComputeWorker) && !n.HasRoles(node.RoleKeyManager) && len(n.Runtimes) > 0 {
-			return fmt.Errorf("registry: sanity check failed: storage worker node shouldn't have any runtimes")
+		if n.HasRoles(node.RoleKeyManager) && len(n.Runtimes) == 0 {
+			return fmt.Errorf("registry: sanity check failed: key manager node must have runtime(s)")
 		}
 
-		if n.HasRoles(node.RoleValidator) && !n.HasRoles(node.RoleComputeWorker) && !n.HasRoles(node.RoleKeyManager) && len(n.Runtimes) > 0 {
+		if n.HasRoles(node.RoleValidator) && !n.HasRoles(node.RoleComputeWorker) && !n.HasRoles(node.RoleStorageWorker) && !n.HasRoles(node.RoleKeyManager) && len(n.Runtimes) > 0 {
 			return fmt.Errorf("registry: sanity check failed: validator node shouldn't have any runtimes")
 		}