INT-1903 added support for allowing to use custom MessageIdGenerationStrategy

spring-integration-core/src/main/java/org/springframework/integration/IntegrationContextRefreshListener.java

@@ -0,0 +1,57 @@
+/*
+ * Copyright 2002-2011 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.integration;
+
+import org.apache.commons.logging.Log;
+import org.apache.commons.logging.LogFactory;
+
+import org.springframework.beans.factory.DisposableBean;
+import org.springframework.beans.factory.NoSuchBeanDefinitionException;
+import org.springframework.context.ApplicationListener;
+import org.springframework.context.event.ContextRefreshedEvent;
+import org.springframework.integration.MessageHeaders.MessageIdGenerationStrategy;
+
+
+/**
+ * @author Oleg Zhurakousky
+ * @since 2.0
+ */
+public class IntegrationContextRefreshListener implements ApplicationListener<ContextRefreshedEvent>, DisposableBean{
+
+	private final Log logger = LogFactory.getLog(getClass());
+	
+	public void onApplicationEvent(ContextRefreshedEvent event) {
+		try {
+			MessageIdGenerationStrategy idGenerationStrategy = 
+				event.getApplicationContext().getBean(MessageIdGenerationStrategy.class);
+			if (logger.isDebugEnabled()) {
+				logger.debug("Using MessageHeaders.MessageIdGenerationStrategy [" + idGenerationStrategy + "]");
+			}
+			MessageHeaders.setMessageIdGenerationStrategy(idGenerationStrategy);
+		}
+		catch (NoSuchBeanDefinitionException ex) {
+			// We need to use the default.
+			if (logger.isDebugEnabled()) {
+				logger.debug("Unable to locate MessageHeaders.MessageIdGenerationStrategy. Will use default UUID.randomUUID()");
+			}
+		}
+	}
+
+	public void destroy() throws Exception {
+		MessageHeaders.reset();
+	}
+
+}

spring-integration-core/src/main/java/org/springframework/integration/MessageHeaders.java

@@ -1,5 +1,5 @@
 /*
- * Copyright 2002-2010 the original author or authors.
+ * Copyright 2002-2011 the original author or authors.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -28,10 +28,14 @@ import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.UUID;
+import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.concurrent.locks.ReentrantReadWriteLock.WriteLock;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
 
+import org.springframework.util.Assert;
+
 /**
  * The headers for a {@link Message}.<br>
  * IMPORTANT: MessageHeaders are immutable. Any mutating operation (e.g., put(..), putAll(..) etc.) 
@@ -50,12 +54,21 @@ import org.apache.commons.logging.LogFactory;
  * 
  * @author Arjen Poutsma
  * @author Mark Fisher
+ * @author Oleg Zhurakousky
  */
 public final class MessageHeaders implements Map<String, Object>, Serializable {
 
 	private static final long serialVersionUID = 6901029029524535147L;
 
 	private static final Log logger = LogFactory.getLog(MessageHeaders.class);
+	
+	private static MessageIdGenerationStrategy messageIdGenerationStrategy = new DefaultIdGenerator();
+	
+	private static final ReentrantReadWriteLock rwl = new ReentrantReadWriteLock();
+
+	private static final WriteLock writeLock = rwl.writeLock();
+	
+	private static boolean idGenerationStrategySet;
 
 	/**
 	 * The key for the Message ID. This is an automatically generated UUID and
@@ -89,11 +102,44 @@ public final class MessageHeaders implements Map<String, Object>, Serializable {
 
 	public MessageHeaders(Map<String, Object> headers) {
 		this.headers = (headers != null) ? new HashMap<String, Object>(headers) : new HashMap<String, Object>();
-		this.headers.put(ID, UUID.randomUUID());
+		/*
+		 * There is a possibility of the race condition when this constructor is called while 
+		 * setMessageIdGenerationStrategy(..) or reset() is invoked, but synchronizing here would be an overkill IMHO. 
+		 * Realistically there will be no messages yet until the ApplicationContext is started 
+		 * (that is when the setMessageIdGenerationStrategy(..) is called and for reset() all the adapters 
+		 * will be shut down by the time reset() is called.
+		 */
+		this.headers.put(ID, MessageHeaders.messageIdGenerationStrategy.generateId());
 		this.headers.put(TIMESTAMP, new Long(System.currentTimeMillis()));
 	}
 
 
+	public static void setMessageIdGenerationStrategy(MessageIdGenerationStrategy messageIdGenerationStrategy) {
+		writeLock.lock();
+		try {
+			Assert.state(!MessageHeaders.idGenerationStrategySet, "'MessageHeaders.messageIdGenerationStrategy' " +
+				"has already been set and can not be set again, unless reset() method is called");
+			logger.info("Message IDs will be generated using custom ID generation strategy: " + messageIdGenerationStrategy);
+			MessageHeaders.messageIdGenerationStrategy = messageIdGenerationStrategy;
+			MessageHeaders.idGenerationStrategySet = true;
+		} 
+		finally {
+			writeLock.unlock();
+		}	
+	}
+	
+	public static void reset(){
+		writeLock.lock();
+		try {
+			MessageHeaders.idGenerationStrategySet = false;
+			MessageHeaders.messageIdGenerationStrategy = new DefaultIdGenerator();
+		} 
+		finally {
+			writeLock.unlock();
+		}
+		logger.info("Message IDs genration strategy was reset to the default");
+	}
+
 	public UUID getId() {
 		return this.get(ID, UUID.class);
 	}
@@ -252,4 +298,16 @@ public final class MessageHeaders implements Map<String, Object>, Serializable {
 		in.defaultReadObject();
 	}
 
+	public static interface MessageIdGenerationStrategy {
+
+		UUID generateId();
+	}
+	
+	private static class DefaultIdGenerator implements MessageIdGenerationStrategy {
+
+		public UUID generateId() {
+			return UUID.randomUUID();
+		}
+		
+	}
 }

spring-integration-core/src/main/java/org/springframework/integration/config/xml/DefaultConfiguringBeanFactoryPostProcessor.java

@@ -20,6 +20,7 @@ import java.util.concurrent.ThreadPoolExecutor.CallerRunsPolicy;
 
 import org.apache.commons.logging.Log;
 import org.apache.commons.logging.LogFactory;
+
 import org.springframework.beans.BeansException;
 import org.springframework.beans.factory.config.BeanDefinition;
 import org.springframework.beans.factory.config.BeanDefinitionHolder;
@@ -53,6 +54,7 @@ class DefaultConfiguringBeanFactoryPostProcessor implements BeanFactoryPostProce
 			this.registerNullChannel(registry);
 			this.registerErrorChannelIfNecessary(registry);
 			this.registerTaskSchedulerIfNecessary(registry);
+			this.registerMessageIdGeneratorIfNecessary(registry);
 		}
 		else if (logger.isWarnEnabled()) {
 			logger.warn("BeanFactory is not a BeanDefinitionRegistry. The default '"
@@ -60,6 +62,11 @@ class DefaultConfiguringBeanFactoryPostProcessor implements BeanFactoryPostProce
 					+ IntegrationContextUtils.TASK_SCHEDULER_BEAN_NAME + "' cannot be configured.");
 		}
 	}
+	
+	private void registerMessageIdGeneratorIfNecessary(BeanDefinitionRegistry registry){
+		String listenerClassName = "org.springframework.integration.IntegrationContextRefreshListener";
+		BeanDefinitionReaderUtils.registerWithGeneratedName(new RootBeanDefinition(listenerClassName), registry);
+	}
 
 	/**
 	 * Register a null channel in the given BeanDefinitionRegistry. The bean name is defined by the constant

spring-integration-core/src/test/java/org/springframework/integration/core/MessageIdGenerationTests-context.xml

@@ -0,0 +1,24 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<beans xmlns="http://www.springframework.org/schema/beans"
+	xmlns:xsi="http://www.w3.org/2001/XMLSchema-instance"
+	xmlns:int="http://www.springframework.org/schema/integration"
+	xsi:schemaLocation="http://www.springframework.org/schema/beans http://www.springframework.org/schema/beans/spring-beans.xsd
+		http://www.springframework.org/schema/integration http://www.springframework.org/schema/integration/spring-integration-2.0.xsd">
+	
+	<bean id="idGenerator" class="org.mockito.Mockito" factory-method="spy">
+		<constructor-arg>
+			<bean class="org.springframework.integration.core.MessageIdGenerationTests.SampleIdGenerator"/>
+		</constructor-arg>
+	</bean>
+	
+	<int:channel id="input"/>
+	
+	<int:service-activator input-channel="input" output-channel="nextA" expression="'nextA'"/>
+	
+	<int:service-activator input-channel="nextA" output-channel="nextB" expression="'nextB'"/>
+	
+	<int:service-activator input-channel="nextB" output-channel="nextC" expression="'nextC'"/>
+	
+	<int:logging-channel-adapter id="nextC" level="WARN" log-full-message="true"/>
+	
+</beans>

spring-integration-core/src/test/java/org/springframework/integration/core/MessageIdGenerationTests.java

@@ -0,0 +1,98 @@
+/*
+ * Copyright 2002-2011 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.integration.core;
+
+import static org.mockito.Mockito.reset;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
+
+import java.util.UUID;
+
+import org.junit.Ignore;
+import org.junit.Test;
+
+import org.springframework.context.support.ClassPathXmlApplicationContext;
+import org.springframework.integration.MessageChannel;
+import org.springframework.integration.MessageHeaders;
+import org.springframework.integration.MessageHeaders.MessageIdGenerationStrategy;
+import org.springframework.integration.message.GenericMessage;
+import org.springframework.util.StopWatch;
+
+/**
+ * @author Oleg Zhurakousky
+ *
+ */
+public class MessageIdGenerationTests {
+
+	@Test
+	public void testCustomIdGeneration(){
+		ClassPathXmlApplicationContext context = new ClassPathXmlApplicationContext("MessageIdGenerationTests-context.xml", this.getClass());
+		MessageIdGenerationStrategy idGenerator = context.getBean("idGenerator", MessageIdGenerationStrategy.class);
+		MessageChannel inputChannel = context.getBean("input", MessageChannel.class);
+		inputChannel.send(new GenericMessage<Integer>(0));
+		verify(idGenerator, times(4)).generateId();
+		reset(idGenerator);
+		context.destroy();
+		new GenericMessage<Integer>(0);
+		verify(idGenerator, times(0)).generateId();
+	}
+	
+	@Test
+	@Ignore
+	public void performanceTest(){
+		int times = 1000000;
+		StopWatch watch = new StopWatch();
+		watch.start();
+		for (int i = 0; i < times; i++) {
+			new GenericMessage<Integer>(0);
+		}
+		watch.stop();
+		double defaultGeneratorElapsedTime = watch.getTotalTimeSeconds();
+		
+		MessageHeaders.setMessageIdGenerationStrategy(new MessageIdGenerationStrategy() {
+			public UUID generateId() {
+				return TimeBasedUUIDGenerator.generateId();
+			}
+		});
+		watch = new StopWatch();
+		watch.start();
+		for (int i = 0; i < times; i++) {
+			new GenericMessage<Integer>(0);
+		}
+		watch.stop();
+		double timebasedGeneratorElapsedTime = watch.getTotalTimeSeconds();
+		
+		System.out.println("Generated " + times + " messages using default UUID generator " +
+				"in " + defaultGeneratorElapsedTime + " seconds");
+		System.out.println("Generated " + times + " messages using Timebased UUID generator " +
+				"in " + timebasedGeneratorElapsedTime + " seconds");
+		
+		System.out.println(defaultGeneratorElapsedTime/timebasedGeneratorElapsedTime);
+	
+	}
+	
+
+	
+	public static class SampleIdGenerator implements MessageIdGenerationStrategy {
+
+		public UUID generateId() {
+			return UUID.nameUUIDFromBytes(((System.currentTimeMillis() - System.nanoTime()) + "").getBytes());
+		}
+		
+	}
+	
+	
+}

spring-integration-core/src/test/java/org/springframework/integration/core/TimeBasedUUIDGenerator.java

@@ -0,0 +1,106 @@
+/*
+ * Copyright 2002-2011 the original author or authors.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.springframework.integration.core;
+
+import java.net.InetAddress;
+import java.net.NetworkInterface;
+import java.util.UUID;
+import java.util.logging.Logger;
+
+
+/**
+ * @author Oleg Zhurakousky
+ *
+ */
+class TimeBasedUUIDGenerator {
+
+	private static final Logger logger = Logger.getLogger(TimeBasedUUIDGenerator.class.getName());
+	
+	public static final Object lock = new Object();
+	
+	private static boolean canNotDetermineMac = true;
+	private static long lastTime;
+	private static long clockSequence = 0;
+	private static final long macAddress = getMac();
+
+	/**
+	 * Will generate unique time based UUID where the next UUID is 
+	 * always greater then the previous.
+	 */
+	public final static UUID generateId() {
+		return generateIdFromTimestamp(System.currentTimeMillis());
+	}
+	
+	public final static UUID generateIdFromTimestamp(long currentTimeMillis){
+		long time;
+		
+		synchronized (lock) {
+			if (currentTimeMillis > lastTime) {
+				lastTime = currentTimeMillis;
+				clockSequence = 0;
+			} else  { 
+				++clockSequence; 
+			}
+		}
+		
+	
+		time = currentTimeMillis;
+		
+		// low Time
+		time = currentTimeMillis << 32;
+		
+		// mid Time
+		time |= ((currentTimeMillis & 0xFFFF00000000L) >> 16);
+
+		// hi Time
+		time |= 0x1000 | ((currentTimeMillis >> 48) & 0x0FFF); // version 1
+		
+		long clock_seq_hi_and_reserved = clockSequence;  
+    	
+    	clock_seq_hi_and_reserved <<=48;	
+  
+    	long cls = 0 | clock_seq_hi_and_reserved;
+    	
+		long lsb = cls | macAddress;
+		if (canNotDetermineMac){
+			logger.warning("UUID generation process was not able to determine your MAC address. Returning random UUID (non version 1 UUID)");
+			return UUID.randomUUID();
+		} else {
+			return new UUID(time, lsb);
+		}
+	}
+	private static final long getMac(){
+		long  macAddressAsLong = 0;
+		try {
+			InetAddress address = InetAddress.getLocalHost();
+			NetworkInterface ni = NetworkInterface.getByInetAddress(address);
+			if (ni != null) {
+				byte[] mac = ni.getHardwareAddress();
+				//Converts array of unsigned bytes to an long
+				if (mac != null) {
+					for (int i = 0; i < mac.length; i++) {					
+						macAddressAsLong <<= 8;
+						macAddressAsLong ^= (long)mac[i] & 0xFF;
+					}
+				}
+			} 
+			canNotDetermineMac = false;
+		} catch (Exception e) {
+			e.printStackTrace();
+		}
+		return macAddressAsLong;
+	}
+}