Merge pull request #185 from benjchristensen/request-context

BugFix: Observable.observeOn Scheduler Lost RequestContext

Author: benjchristensen

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

@@ -73,6 +73,7 @@
 import com.netflix.hystrix.strategy.concurrency.HystrixConcurrencyStrategyDefault;
 import com.netflix.hystrix.strategy.concurrency.HystrixContextCallable;
 import com.netflix.hystrix.strategy.concurrency.HystrixContextRunnable;
+import com.netflix.hystrix.strategy.concurrency.HystrixContextScheduler;
 import com.netflix.hystrix.strategy.concurrency.HystrixRequestContext;
 import com.netflix.hystrix.strategy.eventnotifier.HystrixEventNotifier;
 import com.netflix.hystrix.strategy.executionhook.HystrixCommandExecutionHook;
@@ -848,7 +849,12 @@ public Observable<R> call(Throwable e) {
         });
 
         // we want to hand off work to a different scheduler so we don't tie up the Hystrix thread
-        o = o.observeOn(observeOn);
+        if (!Schedulers.immediate().equals(observeOn)) {
+            // don't waste overhead if it's the 'immediate' scheduler
+            // otherwise we'll 'observeOn' and wrap with the HystrixContextScheduler
+            // to copy state across threads (if threads are involved)
+            o = o.observeOn(new HystrixContextScheduler(observeOn));
+        }
 
         o = o.finallyDo(new Action0() {
 

hystrix-core/src/main/java/com/netflix/hystrix/strategy/concurrency/HystrixConcurrencyStrategy.java

@@ -1,5 +1,5 @@
 /**
- * Copyright 2012 Netflix, Inc.
+ * Copyright 2013 Netflix, Inc.
  * 
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -24,8 +24,17 @@
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicInteger;
 
-import com.netflix.hystrix.HystrixCollapser;
+import org.junit.After;
+import org.junit.Before;
+import org.junit.Test;
+
+import rx.util.functions.Action1;
+import rx.util.functions.Func1;
+
+import com.netflix.config.ConfigurationManager;
 import com.netflix.hystrix.HystrixCommand;
+import com.netflix.hystrix.HystrixCommandGroupKey;
+import com.netflix.hystrix.HystrixRequestLog;
 import com.netflix.hystrix.HystrixThreadPool;
 import com.netflix.hystrix.HystrixThreadPoolKey;
 import com.netflix.hystrix.strategy.HystrixPlugins;
@@ -151,5 +160,64 @@ public void shutdown(T value) {
             };
         };
     }
+    
+    
+    public static class UnitTest {
+        
+        @Before
+        public void prepareForTest() {
+            /* we must call this to simulate a new request lifecycle running and clearing caches */
+            HystrixRequestContext.initializeContext();
+        }
+
+        @After
+        public void cleanup() {
+            // instead of storing the reference from initialize we'll just get the current state and shutdown
+            if (HystrixRequestContext.getContextForCurrentThread() != null) {
+                // it could have been set NULL by the test
+                HystrixRequestContext.getContextForCurrentThread().shutdown();
+            }
+
+            // force properties to be clean as well
+            ConfigurationManager.getConfigInstance().clear();
+        }
+        
+        /**
+         * If the RequestContext does not get transferred across threads correctly this blows up. 
+         * No specific assertions are necessary.
+         */
+        @Test
+        public void testRequestContextPropagatesAcrossObserveOnPool() {
+            new SimpleCommand().execute();
+            new SimpleCommand().observe().map(new Func1<String, String>() {
+
+                @Override
+                public String call(String s) {
+                    System.out.println("Map => Commands: " + HystrixRequestLog.getCurrentRequest().getExecutedCommands());
+                    return s;
+                }
+            }).toBlockingObservable().forEach(new Action1<String>() {
+
+                @Override
+                public void call(String s) {
+                    System.out.println("Result [" + s + "] => Commands: " + HystrixRequestLog.getCurrentRequest().getExecutedCommands());
+                }
+            });
+        }
+        
+        private static class SimpleCommand extends HystrixCommand<String> {
+
+            public SimpleCommand() {
+                super(HystrixCommandGroupKey.Factory.asKey("SimpleCommand"));
+            }
+            
+            @Override
+            protected String run() throws Exception {
+                System.out.println("Executing => Commands: " + HystrixRequestLog.getCurrentRequest().getExecutedCommands());
+                return "Hello";
+            }
+            
+        }
+    }
 
 }

hystrix-core/src/main/java/com/netflix/hystrix/strategy/concurrency/HystrixContextFunc2.java

@@ -0,0 +1,54 @@
+/**
+ * Copyright 2012 Netflix, Inc.
+ * 
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * 
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * 
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.netflix.hystrix.strategy.concurrency;
+
+import rx.Scheduler;
+import rx.Subscription;
+import rx.util.functions.Func2;
+
+/**
+ * Wrapper around {@link Func2} that manages the {@link HystrixRequestContext} initialization and cleanup for the execution of the {@link Func2}
+ * 
+ * @param <T>
+ *            Return type of {@link Func2}
+ * 
+ * @ExcludeFromJavadoc
+ */
+public class HystrixContextFunc2<T> implements Func2<Scheduler, T, Subscription> {
+
+    private final Func2<? super Scheduler, ? super T, ? extends Subscription> actual;
+    private final HystrixRequestContext parentThreadState;
+
+    public HystrixContextFunc2(Func2<? super Scheduler, ? super T, ? extends Subscription> action) {
+        this.actual = action;
+        this.parentThreadState = HystrixRequestContext.getContextForCurrentThread();
+    }
+
+    @Override
+    public Subscription call(Scheduler t1, T t2) {
+        HystrixRequestContext existingState = HystrixRequestContext.getContextForCurrentThread();
+        try {
+            // set the state of this thread to that of its parent
+            HystrixRequestContext.setContextOnCurrentThread(parentThreadState);
+            // execute actual Func2 with the state of the parent
+            return actual.call(t1, t2);
+        } finally {
+            // restore this thread back to its original state
+            HystrixRequestContext.setContextOnCurrentThread(existingState);
+        }
+    }
+
+}

hystrix-core/src/main/java/com/netflix/hystrix/strategy/concurrency/HystrixContextScheduler.java

@@ -0,0 +1,46 @@
+/**
+ * Copyright 2013 Netflix, Inc.
+ * 
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ * 
+ * http://www.apache.org/licenses/LICENSE-2.0
+ * 
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package com.netflix.hystrix.strategy.concurrency;
+
+import java.util.concurrent.TimeUnit;
+
+import rx.Scheduler;
+import rx.Subscription;
+import rx.util.functions.Func2;
+
+/**
+ * Wrap a {@link Scheduler} so that scheduled actions are wrapped with {@link HystrixContextFunc2} so that 
+ * the {@link HystrixRequestContext} is properly copied across threads (if they are used by the {@link Scheduler}).
+ */
+public class HystrixContextScheduler extends Scheduler {
+
+    private final Scheduler actualScheduler;
+
+    public HystrixContextScheduler(Scheduler scheduler) {
+        this.actualScheduler = scheduler;
+    }
+
+    @Override
+    public <T> Subscription schedule(T state, Func2<? super Scheduler, ? super T, ? extends Subscription> action) {
+        return actualScheduler.schedule(state, new HystrixContextFunc2<T>(action));
+    }
+
+    @Override
+    public <T> Subscription schedule(T state, Func2<? super Scheduler, ? super T, ? extends Subscription> action, long delayTime, TimeUnit unit) {
+        return actualScheduler.schedule(state, new HystrixContextFunc2<T>(action), delayTime, unit);
+    }
+
+}