killbill-uncached

Merge pull request #962 from killbill/ro-db-iteration util:

4/30/2018 3:01:11 PM

Changes

Details

diff --git a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/BeatrixIntegrationModule.java b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/BeatrixIntegrationModule.java
index 0fae546..6bdf4e5 100644
--- a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/BeatrixIntegrationModule.java
+++ b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/BeatrixIntegrationModule.java
@@ -24,6 +24,7 @@ import org.killbill.billing.GuicyKillbillTestWithEmbeddedDBModule;
 import org.killbill.billing.account.glue.DefaultAccountModule;
 import org.killbill.billing.api.TestApiListener;
 import org.killbill.billing.beatrix.glue.BeatrixModule;
+import org.killbill.billing.beatrix.integration.db.TestDBRouterAPI;
 import org.killbill.billing.beatrix.integration.overdue.IntegrationTestOverdueModule;
 import org.killbill.billing.beatrix.util.AccountChecker;
 import org.killbill.billing.beatrix.util.AuditChecker;
@@ -57,6 +58,7 @@ import org.killbill.billing.util.glue.ExportModule;
 import org.killbill.billing.util.glue.GlobalLockerModule;
 import org.killbill.billing.util.glue.KillBillModule;
 import org.killbill.billing.util.glue.KillBillShiroModule;
+import org.killbill.billing.util.glue.KillbillApiAopModule;
 import org.killbill.billing.util.glue.NodesModule;
 import org.killbill.billing.util.glue.NonEntityDaoModule;
 import org.killbill.billing.util.glue.RecordIdModule;
@@ -111,6 +113,7 @@ public class BeatrixIntegrationModule extends KillBillModule {
         install(new BroadcastModule(configSource));
         install(new KillBillShiroModuleOnlyIniRealm(configSource));
         install(new BeatrixModule(configSource));
+        install(new KillbillApiAopModule());
 
         bind(AccountChecker.class).asEagerSingleton();
         bind(SubscriptionChecker.class).asEagerSingleton();
@@ -119,6 +122,7 @@ public class BeatrixIntegrationModule extends KillBillModule {
         bind(RefundChecker.class).asEagerSingleton();
         bind(AuditChecker.class).asEagerSingleton();
         bind(TestApiListener.class).asEagerSingleton();
+        bind(TestDBRouterAPI.class).asEagerSingleton();
     }
 
     private final class DefaultInvoiceModuleWithSwitchRepairLogic extends DefaultInvoiceModule {
diff --git a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/db/TestDBRouter.java b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/db/TestDBRouter.java
new file mode 100644
index 0000000..a8d4727
--- /dev/null
+++ b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/db/TestDBRouter.java
@@ -0,0 +1,152 @@
+/*
+ * Copyright 2014-2018 Groupon, Inc
+ * Copyright 2014-2018 The Billing Project, LLC
+ *
+ * The Billing Project licenses this file to you 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 org.killbill.billing.beatrix.integration.db;
+
+import java.util.concurrent.Callable;
+import java.util.concurrent.atomic.AtomicInteger;
+
+import javax.inject.Inject;
+
+import org.joda.time.DateTime;
+import org.killbill.billing.account.api.Account;
+import org.killbill.billing.api.TestApiListener.NextEvent;
+import org.killbill.billing.beatrix.integration.TestIntegrationBase;
+import org.killbill.billing.callcontext.DefaultTenantContext;
+import org.killbill.billing.catalog.api.BillingPeriod;
+import org.killbill.billing.catalog.api.ProductCategory;
+import org.killbill.billing.entitlement.api.DefaultEntitlement;
+import org.killbill.billing.notification.plugin.api.ExtBusEvent;
+import org.killbill.billing.osgi.api.ROTenantContext;
+import org.killbill.billing.util.callcontext.TenantContext;
+import org.killbill.bus.api.PersistentBus.EventBusException;
+import org.testng.annotations.AfterMethod;
+import org.testng.annotations.BeforeMethod;
+import org.testng.annotations.Test;
+
+import com.google.common.eventbus.Subscribe;
+
+import static java.util.concurrent.TimeUnit.SECONDS;
+import static org.awaitility.Awaitility.await;
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+
+public class TestDBRouter extends TestIntegrationBase {
+
+    @Inject
+    private TestDBRouterAPI testDBRouterAPI;
+
+    private PublicListener publicListener;
+    private AtomicInteger externalBusCount;
+
+    @Override
+    @BeforeMethod(groups = "slow")
+    public void beforeMethod() throws Exception {
+        super.beforeMethod();
+
+        this.externalBusCount = new AtomicInteger(0);
+        testDBRouterAPI.reset();
+    }
+
+    @Override
+    protected void registerHandlers() throws EventBusException {
+        super.registerHandlers();
+
+        publicListener = new PublicListener();
+        externalBus.register(publicListener);
+    }
+
+    @AfterMethod(groups = "slow")
+    public void afterMethod() throws Exception {
+        externalBus.unregister(publicListener);
+        super.afterMethod();
+    }
+
+    @Test(groups = "slow")
+    public void testWithBusEvents() throws Exception {
+        final DateTime initialDate = new DateTime(2012, 2, 1, 0, 3, 42, 0, testTimeZone);
+        clock.setTime(initialDate);
+
+        final Account account = createAccountWithNonOsgiPaymentMethod(getAccountData(2));
+        assertNotNull(account);
+
+        final DefaultEntitlement bpEntitlement = createBaseEntitlementAndCheckForCompletion(account.getId(), "externalKey", "Shotgun", ProductCategory.BASE, BillingPeriod.MONTHLY, NextEvent.CREATE, NextEvent.BLOCK, NextEvent.INVOICE);
+        assertNotNull(bpEntitlement);
+
+        await().atMost(10, SECONDS)
+               .until(new Callable<Boolean>() {
+                   @Override
+                   public Boolean call() throws Exception {
+                       // Expecting ACCOUNT_CREATE, ACCOUNT_CHANGE, SUBSCRIPTION_CREATION (2), ENTITLEMENT_CREATE INVOICE_CREATION
+                       return externalBusCount.get() == 6;
+                   }
+               });
+    }
+
+    private void assertNbCalls(final int expectedNbRWCalls, final int expectedNbROCalls) {
+        assertEquals(testDBRouterAPI.getNbRWCalls(), expectedNbRWCalls);
+        assertEquals(testDBRouterAPI.getNbRoCalls(), expectedNbROCalls);
+    }
+
+    public class PublicListener {
+
+        @Subscribe
+        public void handleExternalEvents(final ExtBusEvent event) {
+            testDBRouterAPI.reset();
+
+            final TenantContext tenantContext = new DefaultTenantContext(callContext.getAccountId(), callContext.getTenantId());
+            // Only RO tenant will trigger use of RO DBI (initiated by plugins)
+            final ROTenantContext roTenantContext = new ROTenantContext(tenantContext);
+
+            // RO calls goes to RW DB by default
+            testDBRouterAPI.doROCall(tenantContext);
+            assertNbCalls(1, 0);
+
+            testDBRouterAPI.doROCall(callContext);
+            assertNbCalls(2, 0);
+
+            // Even if the thread is dirty (previous RW calls), the plugin asked for RO DBI
+            testDBRouterAPI.doROCall(roTenantContext);
+            assertNbCalls(2, 1);
+
+            // Make sure subsequent calls go back to the RW DB
+            testDBRouterAPI.doROCall(tenantContext);
+            assertNbCalls(3, 1);
+
+            testDBRouterAPI.doRWCall(callContext);
+            assertNbCalls(4, 1);
+
+            testDBRouterAPI.doROCall(roTenantContext);
+            assertNbCalls(4, 2);
+
+            testDBRouterAPI.doROCall(callContext);
+            assertNbCalls(5, 2);
+
+            testDBRouterAPI.doROCall(tenantContext);
+            assertNbCalls(6, 2);
+
+            testDBRouterAPI.doChainedROCall(tenantContext);
+            assertNbCalls(7, 2);
+
+            testDBRouterAPI.doChainedRWCall(callContext);
+            assertNbCalls(8, 2);
+
+            // Increment only if there are no errors
+            externalBusCount.incrementAndGet();
+        }
+    }
+}
diff --git a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/db/TestDBRouterAPI.java b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/db/TestDBRouterAPI.java
new file mode 100644
index 0000000..73eb7c1
--- /dev/null
+++ b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/db/TestDBRouterAPI.java
@@ -0,0 +1,93 @@
+/*
+ * Copyright 2014-2018 Groupon, Inc
+ * Copyright 2014-2018 The Billing Project, LLC
+ *
+ * The Billing Project licenses this file to you 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 org.killbill.billing.beatrix.integration.db;
+
+import java.util.concurrent.atomic.AtomicInteger;
+
+import javax.inject.Inject;
+
+import org.killbill.billing.KillbillApi;
+import org.killbill.billing.util.callcontext.CallContext;
+import org.killbill.billing.util.callcontext.TenantContext;
+import org.killbill.billing.util.entity.dao.DBRouterUntyped;
+import org.mockito.Mockito;
+import org.mockito.invocation.InvocationOnMock;
+import org.mockito.stubbing.Answer;
+import org.skife.jdbi.v2.Handle;
+import org.skife.jdbi.v2.IDBI;
+
+public class TestDBRouterAPI implements KillbillApi {
+
+    private final AtomicInteger rwCalls = new AtomicInteger(0);
+    private final AtomicInteger roCalls = new AtomicInteger(0);
+
+    private final DBRouterUntyped dbRouter;
+
+    @Inject
+    public TestDBRouterAPI() {
+        final IDBI dbi = Mockito.mock(IDBI.class);
+        Mockito.when(dbi.open()).thenAnswer(new Answer<Handle>() {
+            @Override
+            public Handle answer(final InvocationOnMock invocation) {
+                rwCalls.incrementAndGet();
+                return null;
+            }
+        });
+        final IDBI roDbi = Mockito.mock(IDBI.class);
+        Mockito.when(roDbi.open()).thenAnswer(new Answer<Handle>() {
+            @Override
+            public Handle answer(final InvocationOnMock invocation) {
+                roCalls.incrementAndGet();
+                return null;
+            }
+        });
+
+        this.dbRouter = new DBRouterUntyped(dbi, roDbi);
+    }
+
+    public void reset() {
+        rwCalls.set(0);
+        roCalls.set(0);
+    }
+
+    public void doRWCall(final CallContext callContext) {
+        dbRouter.getHandle(false);
+    }
+
+    public void doROCall(final TenantContext tenantContext) {
+        dbRouter.getHandle(true);
+    }
+
+    // Nesting dolls
+    public void doChainedROCall(final TenantContext tenantContext) {
+        doROCall(tenantContext);
+    }
+
+    // Nesting dolls
+    public void doChainedRWCall(final CallContext callContext) {
+        doRWCall(callContext);
+    }
+
+    public int getNbRWCalls() {
+        return rwCalls.get();
+    }
+
+    public int getNbRoCalls() {
+        return roCalls.get();
+    }
+}
diff --git a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationBase.java b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationBase.java
index 33fe5ab..2c930d4 100644
--- a/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationBase.java
+++ b/beatrix/src/test/java/org/killbill/billing/beatrix/integration/TestIntegrationBase.java
@@ -120,6 +120,7 @@ import org.killbill.billing.util.dao.NonEntityDao;
 import org.killbill.billing.util.nodes.KillbillNodesApi;
 import org.killbill.billing.util.tag.ControlTagType;
 import org.killbill.bus.api.PersistentBus;
+import org.killbill.bus.api.PersistentBus.EventBusException;
 import org.skife.config.ConfigurationObjectFactory;
 import org.skife.config.TimeSpan;
 import org.slf4j.Logger;
@@ -335,12 +336,16 @@ public class TestIntegrationBase extends BeatrixTestSuiteWithEmbeddedDB implemen
 
         // Start services
         lifecycle.fireStartupSequencePriorEventRegistration();
-        busService.getBus().register(busHandler);
+        registerHandlers();
         lifecycle.fireStartupSequencePostEventRegistration();
 
         paymentPlugin.clear();
     }
 
+    protected void registerHandlers() throws EventBusException {
+        busService.getBus().register(busHandler);
+    }
+
     @AfterMethod(groups = "slow")
     public void afterMethod() throws Exception {
         lifecycle.fireShutdownSequencePriorEventUnRegistration();
diff --git a/payment/src/test/java/org/killbill/billing/payment/provider/MockPaymentProviderPlugin.java b/payment/src/test/java/org/killbill/billing/payment/provider/MockPaymentProviderPlugin.java
index b86f476..32bba1c 100644
--- a/payment/src/test/java/org/killbill/billing/payment/provider/MockPaymentProviderPlugin.java
+++ b/payment/src/test/java/org/killbill/billing/payment/provider/MockPaymentProviderPlugin.java
@@ -48,6 +48,8 @@ import org.killbill.billing.util.callcontext.CallContext;
 import org.killbill.billing.util.callcontext.TenantContext;
 import org.killbill.billing.util.entity.DefaultPagination;
 import org.killbill.billing.util.entity.Pagination;
+import org.killbill.billing.util.entity.dao.DBRouterUntyped;
+import org.killbill.billing.util.entity.dao.DBRouterUntyped.THREAD_STATE;
 import org.killbill.clock.Clock;
 
 import com.google.common.base.Preconditions;
@@ -87,6 +89,8 @@ public class MockPaymentProviderPlugin implements PaymentPluginApi {
 
     private final Clock clock;
 
+    private THREAD_STATE lastThreadState = null;
+
     private class InternalPaymentInfo {
 
         private BigDecimal authAmount;
@@ -272,48 +276,60 @@ public class MockPaymentProviderPlugin implements PaymentPluginApi {
         }
     }
 
+    public THREAD_STATE getLastThreadState() {
+        return lastThreadState;
+    }
+
     @Override
     public PaymentTransactionInfoPlugin authorizePayment(final UUID kbAccountId, final UUID kbPaymentId, final UUID kbTransactionId, final UUID kbPaymentMethodId, final BigDecimal amount, final Currency currency, final Iterable<PluginProperty> properties, final CallContext context)
             throws PaymentPluginApiException {
+        updateLastThreadState();
         return getPaymentTransactionInfoPluginResult(kbPaymentId, kbTransactionId, TransactionType.AUTHORIZE, amount, currency, properties);
     }
 
     @Override
     public PaymentTransactionInfoPlugin capturePayment(final UUID kbAccountId, final UUID kbPaymentId, final UUID kbTransactionId, final UUID kbPaymentMethodId, final BigDecimal amount, final Currency currency, final Iterable<PluginProperty> properties, final CallContext context)
             throws PaymentPluginApiException {
+        updateLastThreadState();
         return getPaymentTransactionInfoPluginResult(kbPaymentId, kbTransactionId, TransactionType.CAPTURE, amount, currency, properties);
     }
 
     @Override
     public PaymentTransactionInfoPlugin purchasePayment(final UUID kbAccountId, final UUID kbPaymentId, final UUID kbTransactionId, final UUID kbPaymentMethodId, final BigDecimal amount, final Currency currency, final Iterable<PluginProperty> properties, final CallContext context) throws PaymentPluginApiException {
+        updateLastThreadState();
         return getPaymentTransactionInfoPluginResult(kbPaymentId, kbTransactionId, TransactionType.PURCHASE, amount, currency, properties);
     }
 
     @Override
     public PaymentTransactionInfoPlugin voidPayment(final UUID kbAccountId, final UUID kbPaymentId, final UUID kbTransactionId, final UUID kbPaymentMethodId, final Iterable<PluginProperty> properties, final CallContext context)
             throws PaymentPluginApiException {
+        updateLastThreadState();
         return getPaymentTransactionInfoPluginResult(kbPaymentId, kbTransactionId, TransactionType.VOID, null, null, properties);
     }
 
     @Override
     public PaymentTransactionInfoPlugin creditPayment(final UUID kbAccountId, final UUID kbPaymentId, final UUID kbTransactionId, final UUID kbPaymentMethodId, final BigDecimal amount, final Currency currency, final Iterable<PluginProperty> properties, final CallContext context)
             throws PaymentPluginApiException {
+        updateLastThreadState();
         return getPaymentTransactionInfoPluginResult(kbPaymentId, kbTransactionId, TransactionType.CREDIT, amount, currency, properties);
     }
 
     @Override
     public List<PaymentTransactionInfoPlugin> getPaymentInfo(final UUID kbAccountId, final UUID kbPaymentId, final Iterable<PluginProperty> properties, final TenantContext context) throws PaymentPluginApiException {
+        updateLastThreadState();
         final List<PaymentTransactionInfoPlugin> result = paymentTransactions.get(kbPaymentId.toString());
         return result != null ? result : ImmutableList.<PaymentTransactionInfoPlugin>of();
     }
 
     @Override
     public Pagination<PaymentTransactionInfoPlugin> searchPayments(final String searchKey, final Long offset, final Long limit, final Iterable<PluginProperty> properties, final TenantContext tenantContext) throws PaymentPluginApiException {
+        updateLastThreadState();
         throw new IllegalStateException("Not implemented");
     }
 
     @Override
     public void addPaymentMethod(final UUID kbAccountId, final UUID kbPaymentMethodId, final PaymentMethodPlugin paymentMethodProps, final boolean setDefault, final Iterable<PluginProperty> properties, final CallContext context) throws PaymentPluginApiException {
+        updateLastThreadState();
         // externalPaymentMethodId is set to a random value
         final PaymentMethodPlugin realWithID = new TestPaymentMethodPlugin(kbPaymentMethodId, paymentMethodProps, UUID.randomUUID().toString());
         paymentMethods.put(kbPaymentMethodId.toString(), realWithID);
@@ -324,26 +340,31 @@ public class MockPaymentProviderPlugin implements PaymentPluginApi {
 
     @Override
     public void deletePaymentMethod(final UUID kbAccountId, final UUID kbPaymentMethodId, final Iterable<PluginProperty> properties, final CallContext context) throws PaymentPluginApiException {
+        updateLastThreadState();
         paymentMethods.remove(kbPaymentMethodId.toString());
         paymentMethodsInfo.remove(kbPaymentMethodId.toString());
     }
 
     @Override
     public PaymentMethodPlugin getPaymentMethodDetail(final UUID kbAccountId, final UUID kbPaymentMethodId, final Iterable<PluginProperty> properties, final TenantContext context) throws PaymentPluginApiException {
+        updateLastThreadState();
         return paymentMethods.get(kbPaymentMethodId.toString());
     }
 
     @Override
     public void setDefaultPaymentMethod(final UUID kbAccountId, final UUID kbPaymentMethodId, final Iterable<PluginProperty> properties, final CallContext context) throws PaymentPluginApiException {
+        updateLastThreadState();
     }
 
     @Override
     public List<PaymentMethodInfoPlugin> getPaymentMethods(final UUID kbAccountId, final boolean refreshFromGateway, final Iterable<PluginProperty> properties, final CallContext context) {
+        updateLastThreadState();
         return ImmutableList.<PaymentMethodInfoPlugin>copyOf(paymentMethodsInfo.values());
     }
 
     @Override
     public Pagination<PaymentMethodPlugin> searchPaymentMethods(final String searchKey, final Long offset, final Long limit, final Iterable<PluginProperty> properties, final TenantContext tenantContext) throws PaymentPluginApiException {
+        updateLastThreadState();
         final ImmutableList<PaymentMethodPlugin> results = ImmutableList.<PaymentMethodPlugin>copyOf(Iterables.<PaymentMethodPlugin>filter(paymentMethods.values(), new Predicate<PaymentMethodPlugin>() {
             @Override
             public boolean apply(final PaymentMethodPlugin input) {
@@ -362,6 +383,7 @@ public class MockPaymentProviderPlugin implements PaymentPluginApi {
 
     @Override
     public void resetPaymentMethods(final UUID kbAccountId, final List<PaymentMethodInfoPlugin> input, final Iterable<PluginProperty> properties, final CallContext callContext) {
+        updateLastThreadState();
         paymentMethodsInfo.clear();
         if (input != null) {
             for (final PaymentMethodInfoPlugin cur : input) {
@@ -372,16 +394,19 @@ public class MockPaymentProviderPlugin implements PaymentPluginApi {
 
     @Override
     public HostedPaymentPageFormDescriptor buildFormDescriptor(final UUID kbAccountId, final Iterable<PluginProperty> customFields, final Iterable<PluginProperty> properties, final CallContext callContext) {
+        updateLastThreadState();
         return new DefaultNoOpHostedPaymentPageFormDescriptor(kbAccountId);
     }
 
     @Override
     public GatewayNotification processNotification(final String notification, final Iterable<PluginProperty> properties, final CallContext callContext) throws PaymentPluginApiException {
+        updateLastThreadState();
         return new DefaultNoOpGatewayNotification();
     }
 
     @Override
     public PaymentTransactionInfoPlugin refundPayment(final UUID kbAccountId, final UUID kbPaymentId, final UUID kbTransactionId, final UUID kbPaymentMethodId, final BigDecimal refundAmount, final Currency currency, final Iterable<PluginProperty> properties, final CallContext context) throws PaymentPluginApiException {
+        updateLastThreadState();
 
         final InternalPaymentInfo info = payments.get(kbPaymentId.toString());
         if (info == null) {
@@ -479,4 +504,8 @@ public class MockPaymentProviderPlugin implements PaymentPluginApi {
 
         return result;
     }
+
+    private void updateLastThreadState() {
+        lastThreadState = DBRouterUntyped.getCurrentState();
+    }
 }
diff --git a/payment/src/test/java/org/killbill/billing/payment/TestJanitor.java b/payment/src/test/java/org/killbill/billing/payment/TestJanitor.java
index ecaecf1..1f2e03a 100644
--- a/payment/src/test/java/org/killbill/billing/payment/TestJanitor.java
+++ b/payment/src/test/java/org/killbill/billing/payment/TestJanitor.java
@@ -56,7 +56,10 @@ import org.killbill.billing.payment.provider.DefaultNoOpPaymentInfoPlugin;
 import org.killbill.billing.payment.provider.MockPaymentProviderPlugin;
 import org.killbill.billing.platform.api.KillbillConfigSource;
 import org.killbill.billing.util.callcontext.InternalCallContextFactory;
+import org.killbill.billing.util.entity.dao.DBRouterUntyped;
+import org.killbill.billing.util.entity.dao.DBRouterUntyped.THREAD_STATE;
 import org.killbill.bus.api.PersistentBus.EventBusException;
+import org.killbill.commons.profiling.Profiling.WithProfilingCallback;
 import org.killbill.notificationq.api.NotificationEvent;
 import org.killbill.notificationq.api.NotificationEventWithMetadata;
 import org.killbill.notificationq.api.NotificationQueueService;
@@ -75,8 +78,8 @@ import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
 import com.google.inject.Inject;
 
-import static org.awaitility.Awaitility.await;
 import static java.util.concurrent.TimeUnit.SECONDS;
+import static org.awaitility.Awaitility.await;
 import static org.testng.Assert.assertEquals;
 import static org.testng.Assert.fail;
 
@@ -472,6 +475,45 @@ public class TestJanitor extends PaymentTestSuiteWithEmbeddedDB {
         });
     }
 
+    @Test(groups = "slow")
+    public void testDBRouterThreadState() throws Throwable {
+        final Payment payment = (Payment) DBRouterUntyped.withRODBIAllowed(true,
+                                                                           new WithProfilingCallback<Object, Throwable>() {
+                                                                               @Override
+                                                                               public Payment execute() throws Throwable {
+                                                                                   // Shouldn't happen in practice, but it's just to verify the behavior
+                                                                                   assertEquals(DBRouterUntyped.getCurrentState(), THREAD_STATE.RO_ALLOWED);
+
+                                                                                   final BigDecimal requestedAmount = BigDecimal.TEN;
+                                                                                   testListener.pushExpectedEvent(NextEvent.PAYMENT);
+                                                                                   final Payment payment = paymentApi.createAuthorization(account, account.getPaymentMethodId(), null, requestedAmount, account.getCurrency(), null, UUID.randomUUID().toString(),
+                                                                                                                                          UUID.randomUUID().toString(), ImmutableList.<PluginProperty>of(), callContext);
+                                                                                   testListener.assertListenerStatus();
+
+                                                                                   // Thread switch, RW by default
+                                                                                   assertEquals(mockPaymentProviderPlugin.getLastThreadState(), THREAD_STATE.RW_ONLY);
+                                                                                   // Switched to RW, because of RW DAO call
+                                                                                   assertEquals(DBRouterUntyped.getCurrentState(), THREAD_STATE.RW_ONLY);
+                                                                                   return payment;
+                                                                               }
+                                                                           });
+
+        DBRouterUntyped.withRODBIAllowed(true,
+                                         new WithProfilingCallback<Object, Throwable>() {
+                                             @Override
+                                             public Object execute() throws Throwable {
+                                                 assertEquals(DBRouterUntyped.getCurrentState(), THREAD_STATE.RO_ALLOWED);
+
+                                                 final Payment retrievedPayment2 = paymentApi.getPayment(payment.getId(), true, false, ImmutableList.<PluginProperty>of(), callContext);
+                                                 Assert.assertEquals(retrievedPayment2.getTransactions().get(0).getTransactionStatus(), TransactionStatus.SUCCESS);
+
+                                                 // No thread switch, RO as well
+                                                 assertEquals(mockPaymentProviderPlugin.getLastThreadState(), THREAD_STATE.RO_ALLOWED);
+                                                 assertEquals(DBRouterUntyped.getCurrentState(), THREAD_STATE.RO_ALLOWED);
+                                                 return null;
+                                             }
+                                         });
+    }
 
     private List<PluginProperty> createPropertiesForInvoice(final Invoice invoice) {
         final List<PluginProperty> result = new ArrayList<PluginProperty>();
diff --git a/profiles/killbill/src/main/java/org/killbill/billing/server/modules/JaxRSAopModule.java b/profiles/killbill/src/main/java/org/killbill/billing/server/modules/JaxRSAopModule.java
new file mode 100644
index 0000000..30becf9
--- /dev/null
+++ b/profiles/killbill/src/main/java/org/killbill/billing/server/modules/JaxRSAopModule.java
@@ -0,0 +1,98 @@
+/*
+ * Copyright 2014-2018 Groupon, Inc
+ * Copyright 2014-2018 The Billing Project, LLC
+ *
+ * The Billing Project licenses this file to you 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 org.killbill.billing.server.modules;
+
+import java.lang.reflect.Method;
+
+import javax.ws.rs.DELETE;
+import javax.ws.rs.GET;
+import javax.ws.rs.HEAD;
+import javax.ws.rs.OPTIONS;
+import javax.ws.rs.POST;
+import javax.ws.rs.PUT;
+
+import org.aopalliance.intercept.MethodInterceptor;
+import org.aopalliance.intercept.MethodInvocation;
+import org.killbill.billing.jaxrs.resources.JaxrsResource;
+import org.killbill.billing.util.entity.dao.DBRouterUntyped;
+import org.killbill.billing.util.glue.KillbillApiAopModule;
+import org.killbill.commons.profiling.Profiling.WithProfilingCallback;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import com.google.inject.AbstractModule;
+import com.google.inject.matcher.Matcher;
+import com.google.inject.matcher.Matchers;
+
+public class JaxRSAopModule extends AbstractModule {
+
+    private static final Logger logger = LoggerFactory.getLogger(KillbillApiAopModule.class);
+
+    private static final Matcher<Method> API_RESOURCE_METHOD_MATCHER = new Matcher<Method>() {
+        @Override
+        public boolean matches(final Method method) {
+            return !method.isSynthetic() &&
+                   (
+                           method.getAnnotation(DELETE.class) != null ||
+                           method.getAnnotation(GET.class) != null ||
+                           method.getAnnotation(HEAD.class) != null ||
+                           method.getAnnotation(OPTIONS.class) != null ||
+                           method.getAnnotation(POST.class) != null ||
+                           method.getAnnotation(PUT.class) != null
+                   );
+        }
+
+        @Override
+        public Matcher<Method> and(final Matcher<? super Method> other) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public Matcher<Method> or(final Matcher<? super Method> other) {
+            throw new UnsupportedOperationException();
+        }
+    };
+
+    @Override
+    protected void configure() {
+        bindInterceptor(Matchers.subclassesOf(JaxrsResource.class),
+                        API_RESOURCE_METHOD_MATCHER,
+                        new JaxRsMethodInterceptor());
+    }
+
+    public static class JaxRsMethodInterceptor implements MethodInterceptor {
+
+        @Override
+        public Object invoke(final MethodInvocation invocation) throws Throwable {
+            return DBRouterUntyped.withRODBIAllowed(isRODBIAllowed(invocation),
+                                                    new WithProfilingCallback<Object, Throwable>() {
+                                                        @Override
+                                                        public Object execute() throws Throwable {
+                                                            logger.debug("Entering JAX-RS call {}, arguments: {}", invocation.getMethod(), invocation.getArguments());
+                                                            final Object proceed = invocation.proceed();
+                                                            logger.debug("Exiting  JXA-RS call {}, returning: {}", invocation.getMethod(), proceed);
+                                                            return proceed;
+                                                        }
+                                                    });
+        }
+
+        private boolean isRODBIAllowed(final MethodInvocation invocation) {
+            return invocation.getMethod().getAnnotation(GET.class) != null;
+        }
+    }
+}
diff --git a/profiles/killbill/src/main/java/org/killbill/billing/server/modules/KillbillServerModule.java b/profiles/killbill/src/main/java/org/killbill/billing/server/modules/KillbillServerModule.java
index 6ffbc2a..5d43a8f 100644
--- a/profiles/killbill/src/main/java/org/killbill/billing/server/modules/KillbillServerModule.java
+++ b/profiles/killbill/src/main/java/org/killbill/billing/server/modules/KillbillServerModule.java
@@ -1,7 +1,7 @@
 /*
  * Copyright 2010-2013 Ning, Inc.
- * Copyright 2014-2017 Groupon, Inc
- * Copyright 2014-2017 The Billing Project, LLC
+ * Copyright 2014-2018 Groupon, Inc
+ * Copyright 2014-2018 The Billing Project, LLC
  *
  * The Billing Project licenses this file to you under the Apache License, version 2.0
  * (the "License"); you may not use this file except in compliance with the
@@ -179,6 +179,7 @@ public class KillbillServerModule extends KillbillPlatformModule {
         install(new GlobalLockerModule(configSource));
         install(new KillBillShiroAopModule());
         install(new KillbillApiAopModule());
+        install(new JaxRSAopModule());
         install(new KillBillShiroWebModule(servletContext, skifeConfigSource));
         install(new NonEntityDaoModule(configSource));
         install(new PaymentModule(configSource));
diff --git a/profiles/killbill/src/test/java/org/killbill/billing/jaxrs/resources/TestDBRouterResource.java b/profiles/killbill/src/test/java/org/killbill/billing/jaxrs/resources/TestDBRouterResource.java
new file mode 100644
index 0000000..254d526
--- /dev/null
+++ b/profiles/killbill/src/test/java/org/killbill/billing/jaxrs/resources/TestDBRouterResource.java
@@ -0,0 +1,88 @@
+/*
+ * Copyright 2014-2018 Groupon, Inc
+ * Copyright 2014-2018 The Billing Project, LLC
+ *
+ * The Billing Project licenses this file to you 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 org.killbill.billing.jaxrs.resources;
+
+import java.util.UUID;
+
+import javax.inject.Inject;
+import javax.ws.rs.GET;
+import javax.ws.rs.POST;
+import javax.ws.rs.Path;
+import javax.ws.rs.core.Response;
+
+import org.joda.time.DateTimeZone;
+import org.killbill.billing.GuicyKillbillTestSuite;
+import org.killbill.billing.beatrix.integration.db.TestDBRouterAPI;
+import org.killbill.billing.callcontext.MutableCallContext;
+import org.killbill.billing.callcontext.MutableInternalCallContext;
+import org.killbill.billing.util.callcontext.CallOrigin;
+import org.killbill.billing.util.callcontext.InternalCallContextFactory;
+import org.killbill.billing.util.callcontext.UserType;
+
+@Path("/testDbResource")
+public class TestDBRouterResource implements JaxrsResource {
+
+    private final MutableInternalCallContext internalCallContext = new MutableInternalCallContext(InternalCallContextFactory.INTERNAL_TENANT_RECORD_ID,
+                                                                                                  1687L,
+                                                                                                  DateTimeZone.UTC,
+                                                                                                  GuicyKillbillTestSuite.getClock().getUTCNow(),
+                                                                                                  UUID.randomUUID(),
+                                                                                                  UUID.randomUUID().toString(),
+                                                                                                  CallOrigin.TEST,
+                                                                                                  UserType.TEST,
+                                                                                                  "Testing",
+                                                                                                  "This is a test",
+                                                                                                  GuicyKillbillTestSuite.getClock().getUTCNow(),
+                                                                                                  GuicyKillbillTestSuite.getClock().getUTCNow());
+
+    private final MutableCallContext callContext = new MutableCallContext(internalCallContext);
+
+    private final TestDBRouterAPI testDBRouterAPI;
+
+    @Inject
+    public TestDBRouterResource(final TestDBRouterAPI testDBRouterAPI) {
+        this.testDBRouterAPI = testDBRouterAPI;
+    }
+
+    @POST
+    public Response doChainedRWROCalls() {
+        testDBRouterAPI.reset();
+        testDBRouterAPI.doRWCall(callContext);
+        testDBRouterAPI.doROCall(callContext);
+        return Response.ok().build();
+    }
+
+    @GET
+    public Response doChainedROROCalls() {
+        testDBRouterAPI.reset();
+        testDBRouterAPI.doROCall(callContext);
+        testDBRouterAPI.doROCall(callContext);
+        return Response.ok().build();
+    }
+
+    @GET
+    @Path("/chained")
+    public Response doChainedRORWROCalls() {
+        testDBRouterAPI.reset();
+        testDBRouterAPI.doROCall(callContext);
+        // Naughty: @GET method doing a RW call... Verify the underlying code will detect it and mark the thread as dirty
+        testDBRouterAPI.doRWCall(callContext);
+        testDBRouterAPI.doROCall(callContext);
+        return Response.ok().build();
+    }
+}
diff --git a/profiles/killbill/src/test/java/org/killbill/billing/jaxrs/TestDBRouterResources.java b/profiles/killbill/src/test/java/org/killbill/billing/jaxrs/TestDBRouterResources.java
new file mode 100644
index 0000000..aeabe3d
--- /dev/null
+++ b/profiles/killbill/src/test/java/org/killbill/billing/jaxrs/TestDBRouterResources.java
@@ -0,0 +1,52 @@
+/*
+ * Copyright 2014-2018 Groupon, Inc
+ * Copyright 2014-2018 The Billing Project, LLC
+ *
+ * The Billing Project licenses this file to you 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 org.killbill.billing.jaxrs;
+
+import javax.inject.Inject;
+
+import org.killbill.billing.beatrix.integration.db.TestDBRouterAPI;
+import org.killbill.billing.jaxrs.resources.TestDBRouterResource;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+
+public class TestDBRouterResources extends TestJaxrsBase {
+
+    @Inject
+    private TestDBRouterAPI testDBRouterAPI;
+
+    @Inject
+    private TestDBRouterResource testDBRouterResource;
+
+    @Test(groups = "slow")
+    public void testJaxRSAoPRouting() throws Exception {
+        testDBRouterResource.doChainedROROCalls();
+        assertNbCalls(0, 2);
+
+        testDBRouterResource.doChainedRWROCalls();
+        assertNbCalls(2, 0);
+
+        testDBRouterResource.doChainedRORWROCalls();
+        assertNbCalls(2, 1);
+    }
+
+    private void assertNbCalls(final int expectedNbRWCalls, final int expectedNbROCalls) {
+        assertEquals(testDBRouterAPI.getNbRWCalls(), expectedNbRWCalls);
+        assertEquals(testDBRouterAPI.getNbRoCalls(), expectedNbROCalls);
+    }
+}
diff --git a/profiles/killbill/src/test/java/org/killbill/billing/jaxrs/TestJaxrsBase.java b/profiles/killbill/src/test/java/org/killbill/billing/jaxrs/TestJaxrsBase.java
index 83e3d47..4b07107 100644
--- a/profiles/killbill/src/test/java/org/killbill/billing/jaxrs/TestJaxrsBase.java
+++ b/profiles/killbill/src/test/java/org/killbill/billing/jaxrs/TestJaxrsBase.java
@@ -38,12 +38,14 @@ import org.eclipse.jetty.servlet.FilterHolder;
 import org.joda.time.LocalDate;
 import org.killbill.billing.GuicyKillbillTestWithEmbeddedDBModule;
 import org.killbill.billing.api.TestApiListener;
+import org.killbill.billing.beatrix.integration.db.TestDBRouterAPI;
 import org.killbill.billing.client.KillBillClient;
 import org.killbill.billing.client.KillBillHttpClient;
 import org.killbill.billing.client.model.Payment;
 import org.killbill.billing.client.model.PaymentTransaction;
 import org.killbill.billing.client.model.Tenant;
 import org.killbill.billing.invoice.glue.DefaultInvoiceModule;
+import org.killbill.billing.jaxrs.resources.TestDBRouterResource;
 import org.killbill.billing.jetty.HttpServer;
 import org.killbill.billing.jetty.HttpServerConfig;
 import org.killbill.billing.lifecycle.glue.BusModule;
@@ -75,6 +77,7 @@ import com.google.common.base.Predicate;
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableMap;
 import com.google.common.collect.Iterables;
+import com.google.inject.Binder;
 import com.google.inject.Module;
 import com.google.inject.util.Modules;
 
@@ -146,7 +149,14 @@ public class TestJaxrsBase extends KillbillClient {
         protected Module getModule(final ServletContext servletContext) {
             return Modules.override(new KillbillServerModule(servletContext, serverConfig, configSource)).with(new GuicyKillbillTestWithEmbeddedDBModule(configSource),
                                                                                                                new InvoiceModuleWithMockSender(configSource),
-                                                                                                               new PaymentMockModule(configSource));
+                                                                                                               new PaymentMockModule(configSource),
+                                                                                                               new Module() {
+                                                                                                                   @Override
+                                                                                                                   public void configure(final Binder binder) {
+                                                                                                                       binder.bind(TestDBRouterAPI.class).asEagerSingleton();
+                                                                                                                       binder.bind(TestDBRouterResource.class).asEagerSingleton();
+                                                                                                                   }
+                                                                                                               });
         }
 
     }
diff --git a/profiles/killpay/src/main/java/org/killbill/billing/server/modules/KillpayServerModule.java b/profiles/killpay/src/main/java/org/killbill/billing/server/modules/KillpayServerModule.java
index b4b8f6d..38f2e8e 100644
--- a/profiles/killpay/src/main/java/org/killbill/billing/server/modules/KillpayServerModule.java
+++ b/profiles/killpay/src/main/java/org/killbill/billing/server/modules/KillpayServerModule.java
@@ -1,6 +1,6 @@
 /*
- * Copyright 2014 Groupon, Inc
- * Copyright 2014 The Billing Project, LLC
+ * Copyright 2014-2018 Groupon, Inc
+ * Copyright 2014-2018 The Billing Project, LLC
  *
  * The Billing Project licenses this file to you under the Apache License, version 2.0
  * (the "License"); you may not use this file except in compliance with the
@@ -91,6 +91,7 @@ public class KillpayServerModule extends KillbillServerModule {
         install(new GlobalLockerModule(configSource));
         install(new KillBillShiroAopModule());
         install(new KillbillApiAopModule());
+        install(new JaxRSAopModule());
         install(new KillBillShiroWebModule(servletContext, skifeConfigSource));
         install(new NonEntityDaoModule(configSource));
         install(new PaymentModule(configSource));
diff --git a/subscription/src/test/java/org/killbill/billing/subscription/engine/dao/TestSubscriptionDao.java b/subscription/src/test/java/org/killbill/billing/subscription/engine/dao/TestSubscriptionDao.java
index 253bf4a..5933ced 100644
--- a/subscription/src/test/java/org/killbill/billing/subscription/engine/dao/TestSubscriptionDao.java
+++ b/subscription/src/test/java/org/killbill/billing/subscription/engine/dao/TestSubscriptionDao.java
@@ -40,7 +40,8 @@ import org.killbill.billing.subscription.events.SubscriptionBaseEvent;
 import org.killbill.billing.subscription.events.user.ApiEventBuilder;
 import org.killbill.billing.subscription.events.user.ApiEventCreate;
 import org.killbill.billing.util.UUIDs;
-import org.killbill.billing.util.glue.KillbillApiAopModule;
+import org.killbill.billing.util.entity.dao.DBRouterUntyped;
+import org.killbill.commons.profiling.Profiling.WithProfilingCallback;
 import org.mockito.Mockito;
 import org.skife.jdbi.v2.IDBI;
 import org.testng.Assert;
@@ -175,15 +176,10 @@ public class TestSubscriptionDao extends SubscriptionTestSuiteWithEmbeddedDB {
         assertEquals(result5.get(0).getExternalKey(), bundle2.getExternalKey());
         assertEquals(result5.get(1).getExternalKey(), bundle2.getExternalKey());
         assertEquals(result5.get(2).getExternalKey(), bundle2.getExternalKey());
-
-
     }
 
     @Test(groups = "slow")
-    public void testDirtyFlag() throws SubscriptionBaseApiException {
-        // @BeforeMethod created the account
-        KillbillApiAopModule.resetDirtyDBFlag();
-
+    public void testDirtyFlag() throws Throwable {
         final IDBI dbiSpy = Mockito.spy(dbi);
         final IDBI roDbiSpy = Mockito.spy(roDbi);
         final SubscriptionDao subscriptionDao = new DefaultSubscriptionDao(dbiSpy,
@@ -198,34 +194,49 @@ public class TestSubscriptionDao extends SubscriptionTestSuiteWithEmbeddedDB {
         Mockito.verify(dbiSpy, Mockito.times(0)).open();
         Mockito.verify(roDbiSpy, Mockito.times(0)).open();
 
-        Assert.assertEquals(subscriptionDao.getSubscriptionBundleForAccount(accountId, internalCallContext).size(), 0);
-        Mockito.verify(dbiSpy, Mockito.times(0)).open();
-        Mockito.verify(roDbiSpy, Mockito.times(1)).open();
-
-        Assert.assertEquals(subscriptionDao.getSubscriptionBundleForAccount(accountId, internalCallContext).size(), 0);
-        Mockito.verify(dbiSpy, Mockito.times(0)).open();
-        Mockito.verify(roDbiSpy, Mockito.times(2)).open();
-
-        final String externalKey = UUID.randomUUID().toString();
-        final DateTime startDate = clock.getUTCNow();
-        final DateTime createdDate = startDate.plusSeconds(10);
-        final DefaultSubscriptionBaseBundle bundleDef = new DefaultSubscriptionBaseBundle(externalKey, accountId, startDate, startDate, createdDate, createdDate);
-        final SubscriptionBaseBundle bundle = subscriptionDao.createSubscriptionBundle(bundleDef, catalog, false, internalCallContext);
-        Mockito.verify(dbiSpy, Mockito.times(1)).open();
-        Mockito.verify(roDbiSpy, Mockito.times(2)).open();
-
-        Assert.assertEquals(subscriptionDao.getSubscriptionBundleForAccount(accountId, internalCallContext).size(), 1);
-        Mockito.verify(dbiSpy, Mockito.times(2)).open();
-        Mockito.verify(roDbiSpy, Mockito.times(2)).open();
-
-        Assert.assertEquals(subscriptionDao.getSubscriptionBundleForAccount(accountId, internalCallContext).size(), 1);
-        Mockito.verify(dbiSpy, Mockito.times(3)).open();
-        Mockito.verify(roDbiSpy, Mockito.times(2)).open();
-
-        KillbillApiAopModule.resetDirtyDBFlag();
-
-        Assert.assertEquals(subscriptionDao.getSubscriptionBundleForAccount(accountId, internalCallContext).size(), 1);
-        Mockito.verify(dbiSpy, Mockito.times(3)).open();
-        Mockito.verify(roDbiSpy, Mockito.times(3)).open();
+        // @BeforeMethod created the account
+        DBRouterUntyped.withRODBIAllowed(true,
+                                         new WithProfilingCallback<Object, Throwable>() {
+                                      @Override
+                                      public Object execute() throws Throwable {
+                                          Assert.assertEquals(subscriptionDao.getSubscriptionBundleForAccount(accountId, internalCallContext).size(), 0);
+                                          Mockito.verify(dbiSpy, Mockito.times(0)).open();
+                                          Mockito.verify(roDbiSpy, Mockito.times(1)).open();
+
+                                          Assert.assertEquals(subscriptionDao.getSubscriptionBundleForAccount(accountId, internalCallContext).size(), 0);
+                                          Mockito.verify(dbiSpy, Mockito.times(0)).open();
+                                          Mockito.verify(roDbiSpy, Mockito.times(2)).open();
+
+                                          final String externalKey = UUID.randomUUID().toString();
+                                          final DateTime startDate = clock.getUTCNow();
+                                          final DateTime createdDate = startDate.plusSeconds(10);
+                                          final DefaultSubscriptionBaseBundle bundleDef = new DefaultSubscriptionBaseBundle(externalKey, accountId, startDate, startDate, createdDate, createdDate);
+                                          final SubscriptionBaseBundle bundle = subscriptionDao.createSubscriptionBundle(bundleDef, catalog, false, internalCallContext);
+                                          Mockito.verify(dbiSpy, Mockito.times(1)).open();
+                                          Mockito.verify(roDbiSpy, Mockito.times(2)).open();
+
+                                          Assert.assertEquals(subscriptionDao.getSubscriptionBundleForAccount(accountId, internalCallContext).size(), 1);
+                                          Mockito.verify(dbiSpy, Mockito.times(2)).open();
+                                          Mockito.verify(roDbiSpy, Mockito.times(2)).open();
+
+                                          Assert.assertEquals(subscriptionDao.getSubscriptionBundleForAccount(accountId, internalCallContext).size(), 1);
+                                          Mockito.verify(dbiSpy, Mockito.times(3)).open();
+                                          Mockito.verify(roDbiSpy, Mockito.times(2)).open();
+
+                                          return null;
+                                      }
+                                  });
+
+        DBRouterUntyped.withRODBIAllowed(true,
+                                         new WithProfilingCallback<Object, Throwable>() {
+                                      @Override
+                                      public Object execute() {
+                                          Assert.assertEquals(subscriptionDao.getSubscriptionBundleForAccount(accountId, internalCallContext).size(), 1);
+                                          Mockito.verify(dbiSpy, Mockito.times(3)).open();
+                                          Mockito.verify(roDbiSpy, Mockito.times(3)).open();
+
+                                          return null;
+                                      }
+                                  });
     }
 }
diff --git a/util/src/main/java/org/killbill/billing/util/entity/dao/DBRouterUntyped.java b/util/src/main/java/org/killbill/billing/util/entity/dao/DBRouterUntyped.java
index f369b84..abd2b4b 100644
--- a/util/src/main/java/org/killbill/billing/util/entity/dao/DBRouterUntyped.java
+++ b/util/src/main/java/org/killbill/billing/util/entity/dao/DBRouterUntyped.java
@@ -17,17 +17,29 @@
 
 package org.killbill.billing.util.entity.dao;
 
-import org.killbill.billing.util.glue.KillbillApiAopModule;
+import org.killbill.commons.profiling.Profiling.WithProfilingCallback;
 import org.skife.jdbi.v2.Handle;
 import org.skife.jdbi.v2.IDBI;
 import org.skife.jdbi.v2.TransactionCallback;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
+import com.google.common.annotations.VisibleForTesting;
+
+import static org.killbill.billing.util.entity.dao.DBRouterUntyped.THREAD_STATE.RO_ALLOWED;
+import static org.killbill.billing.util.entity.dao.DBRouterUntyped.THREAD_STATE.RW_ONLY;
+
 public class DBRouterUntyped {
 
     private static final Logger logger = LoggerFactory.getLogger(DBRouterUntyped.class);
 
+    private static final ThreadLocal<THREAD_STATE> CURRENT_THREAD_STATE = new ThreadLocal<THREAD_STATE>() {
+        @Override
+        public THREAD_STATE initialValue() {
+            return RW_ONLY;
+        }
+    };
+
     protected final IDBI dbi;
     protected final IDBI roDbi;
 
@@ -36,6 +48,49 @@ public class DBRouterUntyped {
         this.roDbi = roDbi;
     }
 
+    public static Object withRODBIAllowed(final boolean allowRODBI,
+                                          final WithProfilingCallback<Object, Throwable> callback) throws Throwable {
+        final THREAD_STATE currentState = getCurrentState();
+        CURRENT_THREAD_STATE.set(allowRODBI ? RO_ALLOWED : RW_ONLY);
+
+        try {
+            return callback.execute();
+        } finally {
+            CURRENT_THREAD_STATE.set(currentState);
+        }
+    }
+
+    @VisibleForTesting
+    public static THREAD_STATE getCurrentState() {
+        return CURRENT_THREAD_STATE.get();
+    }
+
+    boolean shouldUseRODBI(final boolean requestedRO) {
+        if (requestedRO) {
+            if (isRODBIAllowed()) {
+                logger.debug("Using RO DBI");
+                return true;
+            } else {
+                // Redirect to the rw instance, to work-around any replication delay
+                logger.debug("RO DBI requested, but thread state is {}, using RW DBI", getCurrentState());
+                return false;
+            }
+        } else {
+            // Disable RO DBI for future calls in this thread
+            disallowRODBI();
+            logger.debug("Using RW DBI");
+            return false;
+        }
+    }
+
+    private boolean isRODBIAllowed() {
+        return getCurrentState() == RO_ALLOWED;
+    }
+
+    private void disallowRODBI() {
+        CURRENT_THREAD_STATE.set(RW_ONLY);
+    }
+
     public Handle getHandle(final boolean requestedRO) {
         if (shouldUseRODBI(requestedRO)) {
             return roDbi.open();
@@ -60,20 +115,10 @@ public class DBRouterUntyped {
         }
     }
 
-    boolean shouldUseRODBI(final boolean requestedRO) {
-        if (!requestedRO) {
-            KillbillApiAopModule.setDirtyDBFlag();
-            logger.debug("Dirty flag set, using RW DBI");
-            return false;
-        } else {
-            if (KillbillApiAopModule.getDirtyDBFlag()) {
-                // Redirect to the rw instance, to work-around any replication delay
-                logger.debug("RO DBI handle requested, but dirty flag set, using RW DBI");
-                return false;
-            } else {
-                logger.debug("Using RO DBI");
-                return true;
-            }
-        }
+    public enum THREAD_STATE {
+        // Advisory that RO DBI can be used
+        RO_ALLOWED,
+        // Dirty flag, calls must go to RW DBI
+        RW_ONLY
     }
 }
diff --git a/util/src/main/java/org/killbill/billing/util/glue/KillbillApiAopModule.java b/util/src/main/java/org/killbill/billing/util/glue/KillbillApiAopModule.java
index 80abe6d..92cc22e 100644
--- a/util/src/main/java/org/killbill/billing/util/glue/KillbillApiAopModule.java
+++ b/util/src/main/java/org/killbill/billing/util/glue/KillbillApiAopModule.java
@@ -22,11 +22,9 @@ import java.lang.reflect.Method;
 import org.aopalliance.intercept.MethodInterceptor;
 import org.aopalliance.intercept.MethodInvocation;
 import org.killbill.billing.KillbillApi;
-import org.killbill.billing.callcontext.DefaultTenantContext;
-import org.killbill.billing.callcontext.InternalCallContext;
-import org.killbill.billing.callcontext.InternalTenantContext;
 import org.killbill.billing.osgi.api.ROTenantContext;
 import org.killbill.billing.util.callcontext.CallContext;
+import org.killbill.billing.util.entity.dao.DBRouterUntyped;
 import org.killbill.commons.profiling.Profiling;
 import org.killbill.commons.profiling.Profiling.WithProfilingCallback;
 import org.killbill.commons.profiling.ProfilingFeature.ProfilingFeatureType;
@@ -40,16 +38,26 @@ import com.google.inject.matcher.Matchers;
 public class KillbillApiAopModule extends AbstractModule {
 
     private static final Logger logger = LoggerFactory.getLogger(KillbillApiAopModule.class);
-    private static final ThreadLocal<Boolean> perThreadDirtyDBFlag = new ThreadLocal<Boolean>();
 
-    static {
-        // Set an initial value
-        resetDirtyDBFlag();
-    }
+    private static final Matcher<Method> SYNTHETIC_METHOD_MATCHER = new Matcher<Method>() {
+        @Override
+        public boolean matches(final Method method) {
+            return method.isSynthetic();
+        }
+
+        @Override
+        public Matcher<Method> and(final Matcher<? super Method> other) {
+            throw new UnsupportedOperationException();
+        }
+
+        @Override
+        public Matcher<Method> or(final Matcher<? super Method> other) {
+            throw new UnsupportedOperationException();
+        }
+    };
 
     @Override
     protected void configure() {
-
         bindInterceptor(Matchers.subclassesOf(KillbillApi.class),
                         Matchers.not(SYNTHETIC_METHOD_MATCHER),
                         new ProfilingMethodInterceptor());
@@ -57,34 +65,41 @@ public class KillbillApiAopModule extends AbstractModule {
 
     public static class ProfilingMethodInterceptor implements MethodInterceptor {
 
-        private final Profiling prof = new Profiling<Object, Throwable>();
+        private final Profiling<Object, Throwable> prof = new Profiling<Object, Throwable>();
 
         @Override
         public Object invoke(final MethodInvocation invocation) throws Throwable {
-            return prof.executeWithProfiling(ProfilingFeatureType.API, invocation.getMethod().getName(), new WithProfilingCallback() {
+            final WithProfilingCallback<Object, Throwable> callback = new WithProfilingCallback<Object, Throwable>() {
                 @Override
                 public Object execute() throws Throwable {
-                    final boolean useRODBIfAvailable = shouldUseRODBIfAvailable(invocation);
-                    if (!useRODBIfAvailable) {
-                        setDirtyDBFlag();
-                    }
-
-                    try {
-                        logger.debug("Entering API call {}, arguments: {}", invocation.getMethod(), invocation.getArguments());
-                        final Object proceed = invocation.proceed();
-                        logger.debug("Exiting  API call {}, returning: {}", invocation.getMethod(), proceed);
-                        return proceed;
-                    } finally {
-                        resetDirtyDBFlag();
-                    }
+                    logger.debug("Entering API call {}, arguments: {}", invocation.getMethod(), invocation.getArguments());
+                    final Object proceed = invocation.proceed();
+                    logger.debug("Exiting  API call {}, returning: {}", invocation.getMethod(), proceed);
+                    return proceed;
                 }
-            });
+            };
+
+            if (forcedRODBI(invocation)) {
+                return prof.executeWithProfiling(ProfilingFeatureType.API,
+                                                 invocation.getMethod().getName(),
+                                                 new WithProfilingCallback<Object, Throwable>() {
+                                                     @Override
+                                                     public Object execute() throws Throwable {
+                                                         return DBRouterUntyped.withRODBIAllowed(true, callback);
+                                                     }
+                                                 });
+            } else {
+                return prof.executeWithProfiling(ProfilingFeatureType.API,
+                                                 invocation.getMethod().getName(),
+                                                 callback);
+            }
         }
 
-        private boolean shouldUseRODBIfAvailable(final MethodInvocation invocation) {
-            // Verify if the flag is already set for re-entrant calls
-            if (getDirtyDBFlag()) {
-                return false;
+        private boolean forcedRODBI(final MethodInvocation invocation) {
+            // Snowflakes from server filters
+            final boolean safeROOperations = "getTenantByApiKey".equals(invocation.getMethod().getName()) || "login".equals(invocation.getMethod().getName());
+            if (safeROOperations) {
+                return true;
             }
 
             final Object[] arguments = invocation.getArguments();
@@ -92,21 +107,11 @@ public class KillbillApiAopModule extends AbstractModule {
                 return false;
             }
 
-            // Snowflakes from server filters
-            final boolean safeROOperations = "getTenantByApiKey".equals(invocation.getMethod().getName()) || "login".equals(invocation.getMethod().getName());
-            if (safeROOperations) {
-                return true;
-            }
-
             for (int i = arguments.length - 1; i >= 0; i--) {
                 final Object argument = arguments[i];
-                // DefaultTenantContext belongs to killbill-internal-api and shouldn't be used by plugins
-                final boolean fromJAXRS = argument instanceof DefaultTenantContext && !(argument instanceof CallContext);
-                // Kill Bill internal re-entrant calls
-                final boolean fromInternalAPIs = argument instanceof InternalTenantContext && !(argument instanceof InternalCallContext);
                 // RO DB explicitly requested by a plugin
                 final boolean pluginRequestROInstance = argument instanceof ROTenantContext && !(argument instanceof CallContext);
-                if (fromJAXRS || fromInternalAPIs || pluginRequestROInstance) {
+                if (pluginRequestROInstance) {
                     return true;
                 }
             }
@@ -114,33 +119,4 @@ public class KillbillApiAopModule extends AbstractModule {
             return false;
         }
     }
-
-    public static void setDirtyDBFlag() {
-        perThreadDirtyDBFlag.set(true);
-    }
-
-    public static void resetDirtyDBFlag() {
-        perThreadDirtyDBFlag.set(false);
-    }
-
-    public static Boolean getDirtyDBFlag() {
-        return perThreadDirtyDBFlag.get() == Boolean.TRUE;
-    }
-
-    private static final Matcher<Method> SYNTHETIC_METHOD_MATCHER = new Matcher<Method>() {
-        @Override
-        public boolean matches(final Method method) {
-            return method.isSynthetic();
-        }
-
-        @Override
-        public Matcher<Method> and(final Matcher<? super Method> other) {
-            throw new UnsupportedOperationException();
-        }
-
-        @Override
-        public Matcher<Method> or(final Matcher<? super Method> other) {
-            throw new UnsupportedOperationException();
-        }
-    };
 }
diff --git a/util/src/test/java/org/killbill/billing/GuicyKillbillTestSuiteWithEmbeddedDB.java b/util/src/test/java/org/killbill/billing/GuicyKillbillTestSuiteWithEmbeddedDB.java
index 37e3f80..719cff7 100644
--- a/util/src/test/java/org/killbill/billing/GuicyKillbillTestSuiteWithEmbeddedDB.java
+++ b/util/src/test/java/org/killbill/billing/GuicyKillbillTestSuiteWithEmbeddedDB.java
@@ -23,7 +23,6 @@ import javax.inject.Named;
 import javax.sql.DataSource;
 
 import org.killbill.billing.util.cache.CacheControllerDispatcher;
-import org.killbill.billing.util.glue.KillbillApiAopModule;
 import org.killbill.commons.embeddeddb.EmbeddedDB;
 import org.skife.jdbi.v2.IDBI;
 import org.slf4j.Logger;
@@ -64,7 +63,6 @@ public class GuicyKillbillTestSuiteWithEmbeddedDB extends GuicyKillbillTestSuite
     public void beforeMethod() throws Exception {
         cleanupAllTables();
         controlCacheDispatcher.clearAll();
-        KillbillApiAopModule.resetDirtyDBFlag();
     }
 
     protected void cleanupAllTables() {