killbill-memoizeit

diff --git a/jaxrs/src/main/java/org/killbill/billing/jaxrs/util/Context.java b/jaxrs/src/main/java/org/killbill/billing/jaxrs/util/Context.java
index b2a67c0..d443a16 100644
--- a/jaxrs/src/main/java/org/killbill/billing/jaxrs/util/Context.java
+++ b/jaxrs/src/main/java/org/killbill/billing/jaxrs/util/Context.java
@@ -1,7 +1,9 @@
 /*
- * Copyright 2010-2013 Ning, Inc.
+ * Copyright 2010-2014 Ning, Inc.
+ * Copyright 2014-2016 Groupon, Inc
+ * Copyright 2014-2016 The Billing Project, LLC
  *
- * Ning licenses this file to you under the Apache License, version 2.0
+ * 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:
  *
@@ -24,6 +26,7 @@ import org.killbill.billing.util.UUIDs;
 import org.killbill.billing.util.callcontext.CallContext;
 import org.killbill.billing.util.callcontext.CallContextFactory;
 import org.killbill.billing.util.callcontext.CallOrigin;
+import org.killbill.billing.util.callcontext.InternalCallContextFactory;
 import org.killbill.billing.util.callcontext.TenantContext;
 import org.killbill.billing.util.callcontext.UserType;
 
@@ -34,13 +37,15 @@ public class Context {
 
     private final CallOrigin origin;
     private final UserType userType;
-    final CallContextFactory contextFactory;
+    private final CallContextFactory contextFactory;
+    private final InternalCallContextFactory internalCallContextFactory;
 
     @Inject
-    public Context(final CallContextFactory factory) {
+    public Context(final CallContextFactory factory, final InternalCallContextFactory internalCallContextFactory) {
         this.origin = CallOrigin.EXTERNAL;
         this.userType = UserType.CUSTOMER;
         this.contextFactory = factory;
+        this.internalCallContextFactory = internalCallContextFactory;
     }
 
     public CallContext createContext(final String createdBy, final String reason, final String comment, final ServletRequest request)
@@ -48,21 +53,31 @@ public class Context {
         try {
             Preconditions.checkNotNull(createdBy, String.format("Header %s needs to be set", JaxrsResource.HDR_CREATED_BY));
             final Tenant tenant = getTenantFromRequest(request);
-            return contextFactory.createCallContext(tenant == null ? null : tenant.getId(), createdBy, origin, userType, reason,
-                                                    comment, UUIDs.randomUUID());
+            final CallContext callContext = contextFactory.createCallContext(tenant == null ? null : tenant.getId(), createdBy, origin, userType, reason,
+                                                                             comment, UUIDs.randomUUID());
+
+            populateMDCContext(callContext);
+
+            return callContext;
         } catch (final NullPointerException e) {
             throw new IllegalArgumentException(e.getMessage());
         }
     }
 
     public TenantContext createContext(final ServletRequest request) {
+        final TenantContext tenantContext;
+
         final Tenant tenant = getTenantFromRequest(request);
         if (tenant == null) {
             // Multi-tenancy may not have been configured - default to "default" tenant (see InternalCallContextFactory)
-            return contextFactory.createTenantContext(null);
+            tenantContext = contextFactory.createTenantContext(null);
         } else {
-            return contextFactory.createTenantContext(tenant.getId());
+            tenantContext = contextFactory.createTenantContext(tenant.getId());
         }
+
+        populateMDCContext(tenantContext);
+
+        return tenantContext;
     }
 
     private Tenant getTenantFromRequest(final ServletRequest request) {
@@ -74,4 +89,9 @@ public class Context {
             return (Tenant) tenantObject;
         }
     }
+
+    private void populateMDCContext(final TenantContext tenantContext) {
+        // InternalCallContextFactory will do it for us
+        internalCallContextFactory.createInternalTenantContext(tenantContext);
+    }
 }

diff --git a/profiles/killbill/src/main/java/org/killbill/billing/server/filters/KillbillMDCInsertingServletFilter.java b/profiles/killbill/src/main/java/org/killbill/billing/server/filters/KillbillMDCInsertingServletFilter.java
new file mode 100644
index 0000000..7782a51
--- /dev/null
+++ b/profiles/killbill/src/main/java/org/killbill/billing/server/filters/KillbillMDCInsertingServletFilter.java
@@ -0,0 +1,83 @@
+/*
+ * Copyright 2016 Groupon, Inc
+ * Copyright 2016 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.filters;
+
+import java.io.IOException;
+import java.io.OutputStream;
+
+import org.killbill.commons.request.Request;
+import org.killbill.commons.request.RequestData;
+import org.slf4j.MDC;
+
+import com.google.inject.Singleton;
+import com.sun.jersey.spi.container.ContainerRequest;
+import com.sun.jersey.spi.container.ContainerRequestFilter;
+import com.sun.jersey.spi.container.ContainerResponse;
+import com.sun.jersey.spi.container.ContainerResponseFilter;
+import com.sun.jersey.spi.container.ContainerResponseWriter;
+
+import static org.killbill.billing.util.callcontext.InternalCallContextFactory.MDC_KB_ACCOUNT_RECORD_ID;
+import static org.killbill.billing.util.callcontext.InternalCallContextFactory.MDC_KB_TENANT_RECORD_ID;
+
+@Singleton
+public class KillbillMDCInsertingServletFilter implements ContainerRequestFilter, ContainerResponseFilter {
+
+    private static final String MDC_REQUEST_ID = "req.requestId";
+
+    @Override
+    public ContainerRequest filter(final ContainerRequest request) {
+        final RequestData perThreadRequestData = Request.getPerThreadRequestData();
+        if (perThreadRequestData != null) {
+            MDC.put(MDC_REQUEST_ID, perThreadRequestData.getRequestId());
+        }
+
+        return request;
+    }
+
+    @Override
+    public ContainerResponse filter(final ContainerRequest request, final ContainerResponse response) {
+        response.setContainerResponseWriter(new Adapter(response.getContainerResponseWriter()));
+        return response;
+    }
+
+    private static final class Adapter implements ContainerResponseWriter {
+
+        private final ContainerResponseWriter crw;
+
+        Adapter(final ContainerResponseWriter containerResponseWriter) {
+            this.crw = containerResponseWriter;
+        }
+
+        @Override
+        public OutputStream writeStatusAndHeaders(final long contentLength, final ContainerResponse response) throws IOException {
+            return crw.writeStatusAndHeaders(contentLength, response);
+        }
+
+        @Override
+        public void finish() throws IOException {
+            crw.finish();
+
+            // Removing possibly inexistent item is OK
+            MDC.remove(MDC_REQUEST_ID);
+
+            // Cleanup
+            MDC.remove(MDC_KB_ACCOUNT_RECORD_ID);
+            MDC.remove(MDC_KB_TENANT_RECORD_ID);
+        }
+    }
+}

diff --git a/profiles/killbill/src/main/java/org/killbill/billing/server/listeners/KillbillGuiceListener.java b/profiles/killbill/src/main/java/org/killbill/billing/server/listeners/KillbillGuiceListener.java
index e205acb..5ec3889 100644
--- a/profiles/killbill/src/main/java/org/killbill/billing/server/listeners/KillbillGuiceListener.java
+++ b/profiles/killbill/src/main/java/org/killbill/billing/server/listeners/KillbillGuiceListener.java
@@ -27,6 +27,7 @@ import org.killbill.billing.jaxrs.resources.JaxRsResourceBase;
 import org.killbill.billing.jaxrs.util.KillbillEventHandler;
 import org.killbill.billing.platform.api.KillbillConfigSource;
 import org.killbill.billing.platform.config.DefaultKillbillConfigSource;
+import org.killbill.billing.server.filters.KillbillMDCInsertingServletFilter;
 import org.killbill.billing.server.filters.ProfilingContainerResponseFilter;
 import org.killbill.billing.server.filters.RequestDataFilter;
 import org.killbill.billing.server.filters.ResponseCorsFilter;
@@ -39,6 +40,7 @@ import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
 import ch.qos.logback.classic.LoggerContext;
+import ch.qos.logback.classic.helpers.MDCInsertingServletFilter;
 import com.google.common.collect.ImmutableMap;
 import com.google.inject.Module;
 import com.google.inject.servlet.ServletModule;
@@ -66,10 +68,11 @@ public class KillbillGuiceListener extends KillbillPlatformGuiceListener {
         // Set the per-thread RequestData first
         builder.addJerseyFilter(RequestDataFilter.class.getName());
 
-        //
-        // Add jersey filters which are executed prior jersey write the output stream
-        //
-        builder.addJerseyFilter("com.sun.jersey.api.container.filter.LoggingFilter");
+        // Logback default MDC
+        builder.addJerseyFilter(MDCInsertingServletFilter.class.getName());
+
+        // Kill Bill specific MDC
+        builder.addJerseyFilter(KillbillMDCInsertingServletFilter.class.getName());
 
         // Disable WADL - it generates noisy log messages, such as:
         // c.s.j.s.w.g.AbstractWadlGeneratorGrammarGenerator - Couldn't find grammar element for class javax.ws.rs.core.Response
@@ -87,10 +90,14 @@ public class KillbillGuiceListener extends KillbillPlatformGuiceListener {
         //builder.addFilter("/" + SWAGGER_PATH + "*", ResponseCorsFilter.class);
         builder.addFilter("/*", ResponseCorsFilter.class);
 
-        // Add TenantFilter right after is multi-tenancy has been configured.
+        // Add TenantFilter right after if multi-tenancy has been configured.
         if (config.isMultiTenancyEnabled()) {
             builder.addFilter("/*", TenantFilter.class);
         }
+
+        // Finally, just before the request starts, enable the LoggingFilter
+        builder.addJerseyFilter("com.sun.jersey.api.container.filter.LoggingFilter");
+
         return builder.build();
     }
 

diff --git a/profiles/killbill/src/main/java/org/killbill/billing/server/security/TenantFilter.java b/profiles/killbill/src/main/java/org/killbill/billing/server/security/TenantFilter.java
index 1a3ad56..8ce5b3e 100644
--- a/profiles/killbill/src/main/java/org/killbill/billing/server/security/TenantFilter.java
+++ b/profiles/killbill/src/main/java/org/killbill/billing/server/security/TenantFilter.java
@@ -37,6 +37,7 @@ import org.apache.shiro.authc.UsernamePasswordToken;
 import org.apache.shiro.authc.pam.ModularRealmAuthenticator;
 import org.apache.shiro.realm.Realm;
 import org.killbill.billing.jaxrs.resources.JaxrsResource;
+import org.killbill.billing.jaxrs.util.Context;
 import org.killbill.billing.server.listeners.KillbillGuiceListener;
 import org.killbill.billing.tenant.api.Tenant;
 import org.killbill.billing.tenant.api.TenantApiException;
@@ -55,6 +56,8 @@ public class TenantFilter implements Filter {
     private static final Logger log = LoggerFactory.getLogger(TenantFilter.class);
 
     @Inject
+    protected Context context;
+    @Inject
     protected TenantUserApi tenantUserApi;
     @Inject
     protected KillbillJdbcTenantRealm killbillJdbcTenantRealm;
@@ -101,6 +104,9 @@ public class TenantFilter implements Filter {
             final Tenant tenant = tenantUserApi.getTenantByApiKey(apiKey);
             request.setAttribute(TENANT, tenant);
 
+            // Create a dummy context, to set the MDC very early for LoggingFilter
+            context.createContext(request);
+
             chain.doFilter(request, response);
         } catch (final TenantApiException e) {
             // Should never happen since Shiro validated the credentials?

diff --git a/profiles/killbill/src/main/resources/logback.xml b/profiles/killbill/src/main/resources/logback.xml
index 9979e27..98d18a0 100644
--- a/profiles/killbill/src/main/resources/logback.xml
+++ b/profiles/killbill/src/main/resources/logback.xml
@@ -18,12 +18,11 @@
 
 <configuration>
     <conversionRule conversionWord="maskedMsg" converterClass="org.killbill.billing.server.log.obfuscators.ObfuscatorConverter" />
-    <conversionRule conversionWord="requestId" converterClass="org.killbill.billing.server.log.RequestIdConverter" />
 
     <appender name="STDOUT" class="ch.qos.logback.core.ConsoleAppender">
         <encoder>
             <!-- See http://jira.qos.ch/browse/LOGBACK-262 -->
-            <pattern>%date{"yyyy-MM-dd'T'HH:mm:ss,SSSZ", UTC} rId='%requestId', th='%thread', lvl='%level', log='%logger{0}', %maskedMsg%n%ex</pattern>
+            <pattern>%date{"yyyy-MM-dd'T'HH:mm:ss,SSSZ", UTC} th='%thread', xff='%X{req.xForwardedFor}', rId='%X{req.requestId}', tRId='%X{kb.tenantRecordId}', aRId='%X{kb.accountRecordId}', lvl='%level', log='%logger{0}', %maskedMsg%n%ex</pattern>
         </encoder>
     </appender>
 

diff --git a/util/src/main/java/org/killbill/billing/util/callcontext/InternalCallContextFactory.java b/util/src/main/java/org/killbill/billing/util/callcontext/InternalCallContextFactory.java
index 0460fd4..c16cd2e 100644
--- a/util/src/main/java/org/killbill/billing/util/callcontext/InternalCallContextFactory.java
+++ b/util/src/main/java/org/killbill/billing/util/callcontext/InternalCallContextFactory.java
@@ -36,6 +36,7 @@ import org.killbill.billing.util.cache.Cachable.CacheType;
 import org.killbill.billing.util.cache.CacheControllerDispatcher;
 import org.killbill.billing.util.dao.NonEntityDao;
 import org.killbill.clock.Clock;
+import org.slf4j.MDC;
 
 import com.google.common.base.MoreObjects;
 
@@ -44,6 +45,9 @@ public class InternalCallContextFactory {
 
     public static final long INTERNAL_TENANT_RECORD_ID = 0L;
 
+    public static final String MDC_KB_ACCOUNT_RECORD_ID = "kb.accountRecordId";
+    public static final String MDC_KB_TENANT_RECORD_ID = "kb.tenantRecordId";
+
     private final ImmutableAccountInternalApi accountInternalApi;
     private final Clock clock;
     private final NonEntityDao nonEntityDao;
@@ -130,6 +134,8 @@ public class InternalCallContextFactory {
      * @return internal tenant callcontext
      */
     public InternalTenantContext createInternalTenantContext(final Long tenantRecordId, @Nullable final Long accountRecordId) {
+        populateMDCContext(accountRecordId, tenantRecordId);
+
         if (accountRecordId == null) {
             return new InternalTenantContext(tenantRecordId);
         } else {
@@ -212,12 +218,18 @@ public class InternalCallContextFactory {
     public InternalCallContext createInternalCallContext(final CallContext context) {
         // If tenant id is null, this will default to the default tenant record id (multi-tenancy disabled)
         final Long tenantRecordId = getTenantRecordIdSafe(context);
+
+        populateMDCContext(null, tenantRecordId);
+
         return new InternalCallContext(tenantRecordId, context);
     }
 
     // Used when we need to re-hydrate the callcontext with the account_record_id (when creating the account)
     public InternalCallContext createInternalCallContext(final Long accountRecordId, final InternalCallContext context) {
         final DateTimeZone accountTimeZone = getAccountTimeZone(context.getTenantRecordId(), accountRecordId);
+
+        populateMDCContext(accountRecordId, context.getTenantRecordId());
+
         return new InternalCallContext(context, accountRecordId, accountTimeZone);
     }
 
@@ -238,15 +250,26 @@ public class InternalCallContextFactory {
         final Long nonNulTenantRecordId = MoreObjects.firstNonNull(tenantRecordId, INTERNAL_TENANT_RECORD_ID);
         final DateTimeZone accountTimeZone = getAccountTimeZone(tenantRecordId, accountRecordId);
 
+        populateMDCContext(accountRecordId, nonNulTenantRecordId);
+
         return new InternalCallContext(nonNulTenantRecordId, accountRecordId, accountTimeZone, userToken, userName, callOrigin, userType, reasonCode, comment,
                                        createdDate, updatedDate);
     }
 
+    private void populateMDCContext(@Nullable final Long accountRecordId, final Long tenantRecordId) {
+        if (accountRecordId != null) {
+            MDC.put(MDC_KB_ACCOUNT_RECORD_ID, String.valueOf(accountRecordId));
+        }
+        MDC.put(MDC_KB_TENANT_RECORD_ID, String.valueOf(tenantRecordId));
+    }
+
     private DateTimeZone getAccountTimeZone(final Long tenantRecordId, @Nullable final Long accountRecordId) {
         if (accountRecordId == null || accountInternalApi == null) {
             return null;
         }
 
+        populateMDCContext(accountRecordId, tenantRecordId);
+
         final InternalTenantContext tmp = new InternalTenantContext(tenantRecordId, accountRecordId, null);
 
         try {