001/**
002 * Licensed to the Apache Software Foundation (ASF) under one or more
003 * contributor license agreements.  See the NOTICE file distributed with
004 * this work for additional information regarding copyright ownership.
005 * The ASF licenses this file to You under the Apache License, Version 2.0
006 * (the "License"); you may not use this file except in compliance with
007 * the License.  You may obtain a copy of the License at
008 *
009 *      http://www.apache.org/licenses/LICENSE-2.0
010 *
011 * Unless required by applicable law or agreed to in writing, software
012 * distributed under the License is distributed on an "AS IS" BASIS,
013 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
014 * See the License for the specific language governing permissions and
015 * limitations under the License.
016 */
017package org.apache.camel.impl;
018
019import org.apache.camel.AsyncCallback;
020import org.apache.camel.Exchange;
021import org.apache.camel.Processor;
022import org.apache.camel.spi.RouteContext;
023import org.apache.camel.spi.UnitOfWork;
024import org.slf4j.Logger;
025import org.slf4j.LoggerFactory;
026import org.slf4j.MDC;
027
028/**
029 * This unit of work supports <a href="http://www.slf4j.org/api/org/slf4j/MDC.html">MDC</a>.
030 *
031 * @version 
032 */
033public class MDCUnitOfWork extends DefaultUnitOfWork {
034
035    public static final String MDC_BREADCRUMB_ID = "camel.breadcrumbId";
036    public static final String MDC_EXCHANGE_ID = "camel.exchangeId";
037    public static final String MDC_MESSAGE_ID = "camel.messageId";
038    public static final String MDC_CORRELATION_ID = "camel.correlationId";
039    public static final String MDC_ROUTE_ID = "camel.routeId";
040    public static final String MDC_CAMEL_CONTEXT_ID = "camel.contextId";
041    public static final String MDC_TRANSACTION_KEY = "camel.transactionKey";
042
043    private static final Logger LOG = LoggerFactory.getLogger(MDCUnitOfWork.class);
044
045    private final String originalBreadcrumbId;
046    private final String originalExchangeId;
047    private final String originalMessageId;
048    private final String originalCorrelationId;
049    private final String originalRouteId;
050    private final String originalCamelContextId;
051    private final String originalTransactionKey;
052
053    public MDCUnitOfWork(Exchange exchange) {
054        super(exchange, LOG);
055
056        // remember existing values
057        this.originalExchangeId = MDC.get(MDC_EXCHANGE_ID);
058        this.originalMessageId = MDC.get(MDC_MESSAGE_ID);
059        this.originalBreadcrumbId = MDC.get(MDC_BREADCRUMB_ID);
060        this.originalCorrelationId = MDC.get(MDC_CORRELATION_ID);
061        this.originalRouteId = MDC.get(MDC_ROUTE_ID);
062        this.originalCamelContextId = MDC.get(MDC_CAMEL_CONTEXT_ID);
063        this.originalTransactionKey = MDC.get(MDC_TRANSACTION_KEY);
064
065        // must add exchange and message id in constructor
066        MDC.put(MDC_EXCHANGE_ID, exchange.getExchangeId());
067        String msgId = exchange.hasOut() ? exchange.getOut().getMessageId() : exchange.getIn().getMessageId();
068        MDC.put(MDC_MESSAGE_ID, msgId);
069        // the camel context id is from exchange
070        MDC.put(MDC_CAMEL_CONTEXT_ID, exchange.getContext().getName());
071        // and add optional correlation id
072        String corrId = exchange.getProperty(Exchange.CORRELATION_ID, String.class);
073        if (corrId != null) {
074            MDC.put(MDC_CORRELATION_ID, corrId);
075        }
076        // and add optional breadcrumb id
077        String breadcrumbId = exchange.getIn().getHeader(Exchange.BREADCRUMB_ID, String.class);
078        if (breadcrumbId != null) {
079            MDC.put(MDC_BREADCRUMB_ID, breadcrumbId);
080        }
081    }
082
083    @Override
084    public UnitOfWork newInstance(Exchange exchange) {
085        return new MDCUnitOfWork(exchange);
086    }
087
088    @Override
089    public void stop() throws Exception {
090        super.stop();
091        // and remove when stopping
092        clear();
093    }
094
095    @Override
096    public void pushRouteContext(RouteContext routeContext) {
097        super.pushRouteContext(routeContext);
098        MDC.put(MDC_ROUTE_ID, routeContext.getRoute().getId());
099    }
100
101    @Override
102    public RouteContext popRouteContext() {
103        RouteContext answer = super.popRouteContext();
104
105        // restore old route id back again after we have popped
106        RouteContext previous = getRouteContext();
107        if (previous != null) {
108            // restore old route id back again
109            MDC.put(MDC_ROUTE_ID, previous.getRoute().getId());
110        } else {
111            // not running in route, so clear (should ideally not happen)
112            MDC.remove(MDC_ROUTE_ID);
113        }
114
115        return answer;
116    }
117
118    @Override
119    public void beginTransactedBy(Object key) {
120        MDC.put(MDC_TRANSACTION_KEY, key.toString());
121        super.beginTransactedBy(key);
122    }
123
124    @Override
125    public void endTransactedBy(Object key) {
126        MDC.remove(MDC_TRANSACTION_KEY);
127        super.endTransactedBy(key);
128    }
129
130    @Override
131    public AsyncCallback beforeProcess(Processor processor, Exchange exchange, AsyncCallback callback) {
132        return new MDCCallback(callback);
133    }
134
135    @Override
136    public void afterProcess(Processor processor, Exchange exchange, AsyncCallback callback, boolean doneSync) {
137        if (!doneSync) {
138            // must clear MDC on current thread as the exchange is being processed asynchronously
139            // by another thread
140            clear();
141        }
142        super.afterProcess(processor, exchange, callback, doneSync);
143    }
144
145    /**
146     * Clears information put on the MDC by this {@link MDCUnitOfWork}
147     */
148    public void clear() {
149        if (this.originalBreadcrumbId != null) {
150            MDC.put(MDC_BREADCRUMB_ID, originalBreadcrumbId);
151        } else {
152            MDC.remove(MDC_BREADCRUMB_ID);
153        }
154        if (this.originalExchangeId != null) {
155            MDC.put(MDC_EXCHANGE_ID, originalExchangeId);
156        } else {
157            MDC.remove(MDC_EXCHANGE_ID);
158        }
159        if (this.originalMessageId != null) {
160            MDC.put(MDC_MESSAGE_ID, originalMessageId);
161        } else {
162            MDC.remove(MDC_MESSAGE_ID);
163        }
164        if (this.originalCorrelationId != null) {
165            MDC.put(MDC_CORRELATION_ID, originalCorrelationId);
166        } else {
167            MDC.remove(MDC_CORRELATION_ID);
168        }
169        if (this.originalRouteId != null) {
170            MDC.put(MDC_ROUTE_ID, originalRouteId);
171        } else {
172            MDC.remove(MDC_ROUTE_ID);
173        }
174        if (this.originalCamelContextId != null) {
175            MDC.put(MDC_CAMEL_CONTEXT_ID, originalCamelContextId);
176        } else {
177            MDC.remove(MDC_CAMEL_CONTEXT_ID);
178        }
179        if (this.originalTransactionKey != null) {
180            MDC.put(MDC_TRANSACTION_KEY, originalTransactionKey);
181        } else {
182            MDC.remove(MDC_TRANSACTION_KEY);
183        }
184    }
185
186    @Override
187    public String toString() {
188        return "MDCUnitOfWork";
189    }
190
191    /**
192     * {@link AsyncCallback} which preserves {@link org.slf4j.MDC} when
193     * the asynchronous routing engine is being used.
194     */
195    private static final class MDCCallback implements AsyncCallback {
196
197        private final AsyncCallback delegate;
198        private final String breadcrumbId;
199        private final String exchangeId;
200        private final String messageId;
201        private final String correlationId;
202        private final String routeId;
203        private final String camelContextId;
204
205        private MDCCallback(AsyncCallback delegate) {
206            this.delegate = delegate;
207            this.exchangeId = MDC.get(MDC_EXCHANGE_ID);
208            this.messageId = MDC.get(MDC_MESSAGE_ID);
209            this.breadcrumbId = MDC.get(MDC_BREADCRUMB_ID);
210            this.correlationId = MDC.get(MDC_CORRELATION_ID);
211            this.camelContextId = MDC.get(MDC_CAMEL_CONTEXT_ID);
212            this.routeId = MDC.get(MDC_ROUTE_ID);
213        }
214
215        public void done(boolean doneSync) {
216            try {
217                if (!doneSync) {
218                    // when done asynchronously then restore information from previous thread
219                    if (breadcrumbId != null) {
220                        MDC.put(MDC_BREADCRUMB_ID, breadcrumbId);
221                    }
222                    if (exchangeId != null) {
223                        MDC.put(MDC_EXCHANGE_ID, exchangeId);
224                    }
225                    if (messageId != null) {
226                        MDC.put(MDC_MESSAGE_ID, messageId);
227                    }
228                    if (correlationId != null) {
229                        MDC.put(MDC_CORRELATION_ID, correlationId);
230                    }
231                    if (camelContextId != null) {
232                        MDC.put(MDC_CAMEL_CONTEXT_ID, camelContextId);
233                    }
234                }
235                // need to setup the routeId finally
236                if (routeId != null) {
237                    MDC.put(MDC_ROUTE_ID, routeId);
238                }
239                
240            } finally {
241                // muse ensure delegate is invoked
242                delegate.done(doneSync);
243            }
244        }
245
246        @Override
247        public String toString() {
248            return delegate.toString();
249        }
250    }
251
252}