[ONOS-7098] Rebalance masters on upgrade completion event

Change-Id: I03e8590f8df57563a3919365e3b382082439e869
diff --git a/core/net/src/main/java/org/onosproject/cluster/impl/MastershipManager.java b/core/net/src/main/java/org/onosproject/cluster/impl/MastershipManager.java
index 40ca1a8..0cab28f 100644
--- a/core/net/src/main/java/org/onosproject/cluster/impl/MastershipManager.java
+++ b/core/net/src/main/java/org/onosproject/cluster/impl/MastershipManager.java
@@ -49,6 +49,9 @@
 import org.onosproject.net.MastershipRole;
 import org.onosproject.net.region.Region;
 import org.onosproject.net.region.RegionService;
+import org.onosproject.upgrade.UpgradeEvent;
+import org.onosproject.upgrade.UpgradeEventListener;
+import org.onosproject.upgrade.UpgradeService;
 import org.slf4j.Logger;
 
 import java.util.ArrayList;
@@ -88,6 +91,7 @@
     private final Logger log = getLogger(getClass());
 
     private final MastershipStoreDelegate delegate = new InternalDelegate();
+    private final UpgradeEventListener upgradeEventListener = new InternalUpgradeEventListener();
 
     @Reference(cardinality = ReferenceCardinality.MANDATORY_UNARY)
     protected MastershipStore store;
@@ -104,13 +108,22 @@
     @Reference(cardinality = ReferenceCardinality.MANDATORY_UNARY)
     protected ComponentConfigService cfgService;
 
+    @Reference(cardinality = ReferenceCardinality.MANDATORY_UNARY)
+    protected UpgradeService upgradeService;
+
     private NodeId localNodeId;
     private Timer requestRoleTimer;
 
     static final boolean DEFAULT_USE_REGION_FOR_BALANCE_ROLES = false;
     @Property(name = "useRegionForBalanceRoles", boolValue = DEFAULT_USE_REGION_FOR_BALANCE_ROLES,
               label = "Use Regions for balancing roles")
-    public boolean useRegionForBalanceRoles;
+    protected boolean useRegionForBalanceRoles;
+
+    private static final boolean DEFAULT_REBALANCE_ROLES_ON_UPGRADE = true;
+    @Property(name = "rebalanceRolesOnUpgrade",
+            boolValue = DEFAULT_REBALANCE_ROLES_ON_UPGRADE,
+            label = "Automatically rebalance roles following an upgrade")
+    protected boolean rebalanceRolesOnUpgrade;
 
     @Activate
     public void activate() {
@@ -492,4 +505,14 @@
         }
     }
 
+    private class InternalUpgradeEventListener implements UpgradeEventListener {
+        @Override
+        public void event(UpgradeEvent event) {
+            if (rebalanceRolesOnUpgrade &&
+                    (event.type() == UpgradeEvent.Type.COMMITTED || event.type() == UpgradeEvent.Type.RESET)) {
+                balanceRoles();
+            }
+        }
+    }
+
 }
diff --git a/core/store/dist/src/main/java/org/onosproject/store/mastership/impl/ConsistentDeviceMastershipStore.java b/core/store/dist/src/main/java/org/onosproject/store/mastership/impl/ConsistentDeviceMastershipStore.java
index 5d9f453..8521f52 100644
--- a/core/store/dist/src/main/java/org/onosproject/store/mastership/impl/ConsistentDeviceMastershipStore.java
+++ b/core/store/dist/src/main/java/org/onosproject/store/mastership/impl/ConsistentDeviceMastershipStore.java
@@ -60,7 +60,6 @@
 import org.onosproject.store.cluster.messaging.MessageSubject;
 import org.onosproject.store.serializers.KryoNamespaces;
 import org.onosproject.store.service.Serializer;
-import org.onosproject.upgrade.UpgradeService;
 import org.slf4j.Logger;
 
 import com.google.common.base.Objects;
@@ -91,9 +90,6 @@
     @Reference(cardinality = ReferenceCardinality.MANDATORY_UNARY)
     protected ClusterCommunicationService clusterCommunicator;
 
-    @Reference(cardinality = ReferenceCardinality.MANDATORY_UNARY)
-    protected UpgradeService upgradeService;
-
     private NodeId localNodeId;
 
     private static final MessageSubject ROLE_RELINQUISH_SUBJECT =