ElastiCacheWithSpringDataRedis.java

/*
 * Copyright 2018 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 example;

import io.lettuce.core.AbstractRedisClient;
import io.lettuce.core.ReadFrom;
import io.lettuce.core.RedisClient;
import io.lettuce.core.RedisURI;
import io.lettuce.core.api.StatefulConnection;
import io.lettuce.core.api.StatefulRedisConnection;
import io.lettuce.core.codec.RedisCodec;
import io.lettuce.core.masterslave.MasterSlave;
import io.lettuce.core.masterslave.StatefulRedisMasterSlaveConnection;
import io.lettuce.core.pubsub.StatefulRedisPubSubConnection;
import io.lettuce.core.sentinel.api.StatefulRedisSentinelConnection;

import java.util.Arrays;
import java.util.List;
import java.util.Optional;

import org.springframework.data.redis.connection.RedisConnection;
import org.springframework.data.redis.connection.RedisStandaloneConfiguration;
import org.springframework.data.redis.connection.lettuce.LettuceClientConfiguration;
import org.springframework.data.redis.connection.lettuce.LettuceConnectionFactory;
import org.springframework.data.redis.connection.lettuce.LettuceConnectionProvider;
import org.springframework.data.redis.connection.lettuce.LettuceConnectionProvider.TargetAware;

/**
 * @author Mark Paluch
 */
public class ElastiCacheWithSpringDataRedis {

	public static void main(String[] args) {

		MyMasterSlaveConfiguration masterSlave = new MyMasterSlaveConfiguration(
				Arrays.asList(RedisURI.create("127.0.0.1", 6379), RedisURI.create("127.0.0.1", 6380)));
		LettuceClientConfiguration configuration = LettuceClientConfiguration.builder().readFrom(ReadFrom.SLAVE).build();

		MyLettuceConnectionFactory factory = new MyLettuceConnectionFactory(masterSlave, configuration);

		factory.afterPropertiesSet();

		RedisConnection connection = factory.getConnection();

		connection.set("foo".getBytes(), "bar".getBytes());

		System.out.println(new String(connection.get("foo".getBytes())));

		connection.close();
	}

	static class MyLettuceConnectionFactory extends LettuceConnectionFactory {

		private final MyMasterSlaveConfiguration configuration;

		public MyLettuceConnectionFactory(MyMasterSlaveConfiguration standaloneConfig,
				LettuceClientConfiguration clientConfig) {
			super(standaloneConfig, clientConfig);
			this.configuration = standaloneConfig;
		}

		@Override
		protected LettuceConnectionProvider doCreateConnectionProvider(AbstractRedisClient client, RedisCodec<?, ?> codec) {
			return new ElasticacheConnectionProvider((RedisClient) client, codec, getClientConfiguration().getReadFrom(),
					this.configuration);
		}
	}

	static class MyMasterSlaveConfiguration extends RedisStandaloneConfiguration {

		private final List<RedisURI> endpoints;

		public MyMasterSlaveConfiguration(List<RedisURI> endpoints) {
			this.endpoints = endpoints;
		}

		public List<RedisURI> getEndpoints() {
			return endpoints;
		}
	}

	static class ElasticacheConnectionProvider implements LettuceConnectionProvider, TargetAware {

		private final RedisClient client;
		private final RedisCodec<?, ?> codec;
		private final Optional<ReadFrom> readFrom;
		private final MyMasterSlaveConfiguration configuration;

		public ElasticacheConnectionProvider(RedisClient client, RedisCodec<?, ?> codec, Optional<ReadFrom> readFrom,
				MyMasterSlaveConfiguration configuration) {

			this.client = client;
			this.codec = codec;
			this.readFrom = readFrom;
			this.configuration = configuration;
		}

		@Override
		public <T extends StatefulConnection<?, ?>> T getConnection(Class<T> connectionType) {

			if (connectionType.equals(StatefulRedisSentinelConnection.class)) {
				return connectionType.cast(client.connectSentinel());
			}

			if (connectionType.equals(StatefulRedisPubSubConnection.class)) {
				return connectionType.cast(client.connectPubSub(codec));
			}

			if (StatefulConnection.class.isAssignableFrom(connectionType)) {

				return connectionType.cast(readFrom.map(it -> this.masterSlaveConnection(configuration.getEndpoints(), it))
						.orElseGet(() -> client.connect(codec)));
			}

			throw new UnsupportedOperationException("Connection type " + connectionType + " not supported!");
		}

		@Override
		public <T extends StatefulConnection<?, ?>> T getConnection(Class<T> connectionType, RedisURI redisURI) {

			if (connectionType.equals(StatefulRedisSentinelConnection.class)) {
				return connectionType.cast(client.connectSentinel(redisURI));
			}

			if (connectionType.equals(StatefulRedisPubSubConnection.class)) {
				return connectionType.cast(client.connectPubSub(codec, redisURI));
			}

			if (StatefulConnection.class.isAssignableFrom(connectionType)) {

				if (readFrom.isPresent()) {
					throw new UnsupportedOperationException("Not supported. Configured to Master/Slave with multiple URLs");
				}
				return connectionType.cast(client.connect(codec));
			}

			throw new UnsupportedOperationException("Connection type " + connectionType + " not supported!");
		}

		private StatefulRedisConnection masterSlaveConnection(List<RedisURI> endpoints, ReadFrom readFrom) {

			StatefulRedisMasterSlaveConnection<?, ?> connection = MasterSlave.connect(client, codec, endpoints);
			connection.setReadFrom(readFrom);

			return connection;
		}
	}
}