Fixing the resettability of HystrixMetricsPublisherFactory

hystrix-core/src/main/java/com/netflix/hystrix/Hystrix.java

@@ -5,6 +5,7 @@
 import java.util.concurrent.TimeUnit;
 
 import com.netflix.hystrix.strategy.HystrixPlugins;
+import com.netflix.hystrix.strategy.metrics.HystrixMetricsPublisherFactory;
 import com.netflix.hystrix.strategy.properties.HystrixPropertiesFactory;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

hystrix-core/src/main/java/com/netflix/hystrix/strategy/HystrixPlugins.java

@@ -25,6 +25,7 @@
 import com.netflix.hystrix.strategy.executionhook.HystrixCommandExecutionHookDefault;
 import com.netflix.hystrix.strategy.metrics.HystrixMetricsPublisher;
 import com.netflix.hystrix.strategy.metrics.HystrixMetricsPublisherDefault;
+import com.netflix.hystrix.strategy.metrics.HystrixMetricsPublisherFactory;
 import com.netflix.hystrix.strategy.properties.HystrixPropertiesStrategy;
 import com.netflix.hystrix.strategy.properties.HystrixPropertiesStrategyDefault;
 
@@ -64,6 +65,7 @@ public static void reset() {
         getInstance().metricsPublisher.set(null);
         getInstance().propertiesFactory.set(null);
         getInstance().commandExecutionHook.set(null);
+        HystrixMetricsPublisherFactory.reset();
     }
 
     /**

hystrix-core/src/main/java/com/netflix/hystrix/strategy/metrics/HystrixMetricsPublisherFactory.java

@@ -16,6 +16,7 @@
 package com.netflix.hystrix.strategy.metrics;
 
 import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.atomic.AtomicReference;
 
 import com.netflix.hystrix.HystrixCircuitBreaker;
 import com.netflix.hystrix.HystrixCommand;
@@ -81,30 +82,16 @@ public static HystrixMetricsPublisherCommand createOrRetrievePublisherForCommand
 
     /**
      * Resets the SINGLETON object.
-     *
-     */
-    /* package */ static void reset() {
-        SINGLETON = new HystrixMetricsPublisherFactory();
-    }
-
-    /**
      * Clears all state from publishers. If new requests come in instances will be recreated.
      *
      */
-    /* package */ void _reset() {
-        commandPublishers.clear();
-        threadPoolPublishers.clear();
-    }
-
-    private final HystrixMetricsPublisher strategy;
-
-    private HystrixMetricsPublisherFactory() {
-        this(HystrixPlugins.getInstance().getMetricsPublisher());
+    public static void reset() {
+        SINGLETON = new HystrixMetricsPublisherFactory();
+        SINGLETON.commandPublishers.clear();
+        SINGLETON.threadPoolPublishers.clear();
     }
 
-    /* package */ HystrixMetricsPublisherFactory(HystrixMetricsPublisher strategy) {
-        this.strategy = strategy;
-    }
+    /* package */ HystrixMetricsPublisherFactory()  {}
 
     // String is CommandKey.name() (we can't use CommandKey directly as we can't guarantee it implements hashcode/equals correctly)
     private final ConcurrentHashMap<String, HystrixMetricsPublisherCommand> commandPublishers = new ConcurrentHashMap<String, HystrixMetricsPublisherCommand>();
@@ -116,7 +103,7 @@ private HystrixMetricsPublisherFactory() {
             return publisher;
         }
         // it doesn't exist so we need to create it
-        publisher = strategy.getMetricsPublisherForCommand(commandKey, commandOwner, metrics, circuitBreaker, properties);
+        publisher = HystrixPlugins.getInstance().getMetricsPublisher().getMetricsPublisherForCommand(commandKey, commandOwner, metrics, circuitBreaker, properties);
         // attempt to store it (race other threads)
         HystrixMetricsPublisherCommand existing = commandPublishers.putIfAbsent(commandKey.name(), publisher);
         if (existing == null) {
@@ -141,7 +128,7 @@ private HystrixMetricsPublisherFactory() {
             return publisher;
         }
         // it doesn't exist so we need to create it
-        publisher = strategy.getMetricsPublisherForThreadPool(threadPoolKey, metrics, properties);
+        publisher = HystrixPlugins.getInstance().getMetricsPublisher().getMetricsPublisherForThreadPool(threadPoolKey, metrics, properties);
         // attempt to store it (race other threads)
         HystrixMetricsPublisherThreadPool existing = threadPoolPublishers.putIfAbsent(threadPoolKey.name(), publisher);
         if (existing == null) {

hystrix-core/src/test/java/com/netflix/hystrix/HystrixThreadPoolTest.java

@@ -79,14 +79,12 @@ public HystrixThreadPoolMetrics getHystrixThreadPoolMetrics() {
 
     @Test
     public void ensureThreadPoolInstanceIsTheOneRegisteredWithMetricsPublisherAndThreadPoolCache() throws IllegalAccessException, NoSuchFieldException {
-        new HystrixPluginsTest().reset();
         HystrixPlugins.getInstance().registerMetricsPublisher(new HystrixMetricsPublisher() {
             @Override
             public HystrixMetricsPublisherThreadPool getMetricsPublisherForThreadPool(HystrixThreadPoolKey threadPoolKey, HystrixThreadPoolMetrics metrics, HystrixThreadPoolProperties properties) {
                 return new HystrixMetricsPublisherThreadPoolContainer(metrics);
             }
         });
-        new HystrixMetricsPublisherFactoryTest().reset();
         HystrixThreadPoolKey threadPoolKey = HystrixThreadPoolKey.Factory.asKey("threadPoolFactoryConcurrencyTest");
         HystrixThreadPool poolOne = new HystrixThreadPool.HystrixThreadPoolDefault(
                 threadPoolKey, HystrixThreadPoolProperties.Setter.getUnitTestPropertiesBuilder());

hystrix-core/src/test/java/com/netflix/hystrix/strategy/HystrixPluginsTest.java

@@ -29,12 +29,7 @@
 public class HystrixPluginsTest {
     @After
     public void reset() {
-        // use private access to reset so we can test different initializations via the public static flow
-        HystrixPlugins.getInstance().concurrencyStrategy.set(null);
-        HystrixPlugins.getInstance().metricsPublisher.set(null);
-        HystrixPlugins.getInstance().notifier.set(null);
-        HystrixPlugins.getInstance().propertiesFactory.set(null);
-        HystrixPlugins.getInstance().commandExecutionHook.set(null);
+        HystrixPlugins.reset();
     }
 
     @Test
@@ -241,5 +236,4 @@ protected Void run() throws Exception {
             return null;
         }
     }
-
 }

hystrix-core/src/test/java/com/netflix/hystrix/strategy/metrics/HystrixMetricsPublisherFactoryTest.java

@@ -1,10 +1,13 @@
 package com.netflix.hystrix.strategy.metrics;
 
+import static junit.framework.Assert.assertNotSame;
 import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertSame;
 
 import java.util.ArrayList;
 import java.util.concurrent.atomic.AtomicInteger;
 
+import com.netflix.hystrix.strategy.HystrixPlugins;
 import org.junit.Before;
 import org.junit.Test;
 
@@ -20,7 +23,7 @@
 public class HystrixMetricsPublisherFactoryTest {
     @Before
     public void reset() {
-        HystrixMetricsPublisherFactory.reset();
+        HystrixPlugins.reset();
     }
 
     /**
@@ -29,7 +32,8 @@ public void reset() {
     @Test
     public void testSingleInitializePerKey() {
         final TestHystrixMetricsPublisher publisher = new TestHystrixMetricsPublisher();
-        final HystrixMetricsPublisherFactory factory = new HystrixMetricsPublisherFactory(publisher);
+        HystrixPlugins.getInstance().registerMetricsPublisher(publisher);
+        final HystrixMetricsPublisherFactory factory = new HystrixMetricsPublisherFactory();
         ArrayList<Thread> threads = new ArrayList<Thread>();
         for (int i = 0; i < 20; i++) {
             threads.add(new Thread(new Runnable() {
@@ -63,6 +67,32 @@ public void run() {
         assertEquals(1, publisher.threadCounter.get());
     }
 
+    @Test
+    public void testMetricsPublisherReset() {
+        // precondition: HystrixMetricsPublisherFactory class is not loaded. Calling HystrixPlugins.reset() here should be good enough to run this with other tests.
+
+        // set first custom publisher
+        HystrixCommandKey key = HystrixCommandKey.Factory.asKey("key");
+        HystrixMetricsPublisherCommand firstCommand = new HystrixMetricsPublisherCommandDefault(key, null, null, null, null);
+        HystrixMetricsPublisher firstPublisher = new CustomPublisher(firstCommand);
+        HystrixPlugins.getInstance().registerMetricsPublisher(firstPublisher);
+
+        // ensure that first custom publisher is used
+        HystrixMetricsPublisherCommand cmd = HystrixMetricsPublisherFactory.createOrRetrievePublisherForCommand(key, null, null, null, null);
+        assertSame(firstCommand, cmd);
+
+        // reset, then change to second custom publisher
+        HystrixPlugins.reset();
+        HystrixMetricsPublisherCommand secondCommand = new HystrixMetricsPublisherCommandDefault(key, null, null, null, null);
+        HystrixMetricsPublisher secondPublisher = new CustomPublisher(secondCommand);
+        HystrixPlugins.getInstance().registerMetricsPublisher(secondPublisher);
+
+        // ensure that second custom publisher is used
+        cmd = HystrixMetricsPublisherFactory.createOrRetrievePublisherForCommand(key, null, null, null, null);
+        assertNotSame(firstCommand, cmd);
+        assertSame(secondCommand, cmd);
+    }
+
     private static class TestHystrixMetricsPublisher extends HystrixMetricsPublisher {
 
         AtomicInteger commandCounter = new AtomicInteger();
@@ -97,4 +127,16 @@ private static enum TestCommandKey implements HystrixCommandKey {
     private static enum TestThreadPoolKey implements HystrixThreadPoolKey {
         TEST_A, TEST_B;
     }
+
+    static class CustomPublisher extends HystrixMetricsPublisher{
+        private HystrixMetricsPublisherCommand commandToReturn;
+        public CustomPublisher(HystrixMetricsPublisherCommand commandToReturn){
+            this.commandToReturn = commandToReturn;
+        }
+
+        @Override
+        public HystrixMetricsPublisherCommand getMetricsPublisherForCommand(HystrixCommandKey commandKey, HystrixCommandGroupKey commandGroupKey, HystrixCommandMetrics metrics, HystrixCircuitBreaker circuitBreaker, HystrixCommandProperties properties) {
+            return commandToReturn;
+        }
+    }
 }