From 50d1c7b4816baefe4b47fd59271001d5d590cd3f Mon Sep 17 00:00:00 2001
From: Shalin Shekhar Mangar <shalin@apache.org>
Date: Wed, 10 Oct 2018 17:12:50 +0530
Subject: [PATCH] SOLR-12739: Release the policy session as soon as we're done
 with the computation.

This fixes the CollectionsAPIDistributedZkTest.testCoresAreDistributedAcrossNodes test failures. Due to the various tests for exceptional conditions, there were times where the session was not released causing stale data to remain in the policy session cache.
---
 .../cloud/api/collections/AddReplicaCmd.java  | 19 +++++++++++--------
 1 file changed, 11 insertions(+), 8 deletions(-)

diff --git a/solr/core/src/java/org/apache/solr/cloud/api/collections/AddReplicaCmd.java b/solr/core/src/java/org/apache/solr/cloud/api/collections/AddReplicaCmd.java
index 6e851dbd957..8b72cdf2923 100644
--- a/solr/core/src/java/org/apache/solr/cloud/api/collections/AddReplicaCmd.java
+++ b/solr/core/src/java/org/apache/solr/cloud/api/collections/AddReplicaCmd.java
@@ -141,10 +141,17 @@ public class AddReplicaCmd implements OverseerCollectionMessageHandler.Cmd {
     }
 
     AtomicReference<PolicyHelper.SessionWrapper> sessionWrapper = new AtomicReference<>();
-    List<CreateReplica> createReplicas = buildReplicaPositions(ocmh.cloudManager, clusterState, collectionName, message, replicaTypesVsCount, sessionWrapper)
-        .stream()
-        .map(replicaPosition -> assignReplicaDetails(ocmh.cloudManager, clusterState, message, replicaPosition))
-        .collect(Collectors.toList());
+    List<CreateReplica> createReplicas;
+    try {
+      createReplicas = buildReplicaPositions(ocmh.cloudManager, clusterState, collectionName, message, replicaTypesVsCount, sessionWrapper)
+          .stream()
+          .map(replicaPosition -> assignReplicaDetails(ocmh.cloudManager, clusterState, message, replicaPosition))
+          .collect(Collectors.toList());
+    } finally {
+      if (sessionWrapper.get() != null) {
+        sessionWrapper.get().release();
+      }
+    }
 
     ShardHandler shardHandler = ocmh.shardHandlerFactory.getShardHandler();
     ZkStateReader zkStateReader = ocmh.zkStateReader;
@@ -162,10 +169,6 @@ public class AddReplicaCmd implements OverseerCollectionMessageHandler.Cmd {
       for (CreateReplica replica : createReplicas) {
         ocmh.waitForCoreNodeName(collectionName, replica.node, replica.coreName);
       }
-
-      if (sessionWrapper.get() != null) {
-        sessionWrapper.get().release();
-      }
       if (onComplete != null) onComplete.run();
     };