DB분산처리를 위한 sharding

public enum ShardingStrategy implements IdentityComparable<ShardingStrategy> { // [1]
    RANGE, MODULAR
}

@Getter
@Setter
public class ShardingDataSourceProperty {
    private List<Shard> shards;

    @Getter
    @Setter
    public static class Shard {
        private String username;
        private String password;
        private Property master;
        private List<Property> slaves;
    }

    @Getter
    @Setter
    public static class Property {
        private String name;
        private String url;
    }
}

@Getter
@Setter
public class ShardingProperty {
    private ShardingStrategy strategy;
    private List<ShardingRule> rules;
    private int mod;

    @Getter
    @Setter
    public static class ShardingRule {
        private int shardNo;
        private long rangeMin;
        private long rangeMax;
    }
}

@Configuration
@EnableJpaRepositories( ... )
@ConfigurationProperties(prefix = "datasource")
public class FriendConfig {

    private ShardingDataSourceProperty friend; // [1]

    @Bean
    public DataSource friendDataSource() {
        DataSourceRouter router = new DataSourceRouter(); // [2]
        Map<Object, Object> dataSourceMap = new LinkedHashMap<>(); // [3]

        for (int i = 0; i < property.getShards().size(); i++) {
            ShardingDataSourceProperty.Shard shard = property.getShards().get(i);

            DataSource masterDs = dataSource(shard.getUsername(), shard.getPassword(), shard.getMaster().getUrl());
            dataSourceMap.put(i + SHARD_DELIMITER + shard.getMaster().getName(), masterDs); // [4]

            for (ShardingDataSourceProperty.Property slave : shard.getSlaves()) {
                DataSource slaveDs = dataSource(shard.getUsername(), shard.getPassword(), slave.getUrl());
                dataSourceMap.put(i + SHARD_DELIMITER + slave.getName(), slaveDs);
            }
        }

        router.setTargetDataSources(dataSourceMap); // [5]
        router.afterPropertiesSet(); // [6]

        return new LazyConnectionDataSourceProxy(router); // [7]
    }

    ...
}

public enum ShardingTarget implements IdentityComparable<ShardingTarget> {
    FRIEND, ...
}

@Setter
public class ShardingConfig {
    private static Map<ShardingTarget, ShardingProperty> shardingPropertyMap = new ConcurrentHashMap<>(); // [1]
    
    ...

    public static Map<ShardingTarget, ShardingProperty> getShardingPropertyMap() { 
        return shardingPropertyMap;
    }
}

@Configuration
@ConfigurationProperties(prefix = "sharding")
@Setter
public class FriendShardingConfig {

    private ShardingProperty friend;

    @PostConstruct
    public void init() {
        ShardingConfig.getShardingPropertyMap().put(ShardingTarget.FRIEND, friend); // [1]
    }
}

public class UserHolder {

    private static final ThreadLocal<Context> userContext = new ThreadLocal<>(); // [1]

    ...
    public static void setSharding(ShardingTarget target, long shardKey) { // [2]
        getUserContext().setSharding(new Sharding(target, shardKey));
    }

    public static void clearSharding() { // [3]
        getUserContext().setSharding(null);
    }

    public static Sharding getSharding() {
        return getUserContext().getSharding();
    }
    ...

    @Getter
    @Setter
    public static class Context {
        ...
        private Sharding sharding;
    }

    @Getter
    @Setter
    public static class Sharding { // [4]
        private ShardingTarget target;
        private long shardKey;

        Sharding(ShardingTarget target, long shardKey) {
            this.target = target;
            this.shardKey = shardKey;
        }
    }

public class DataSourceRouter extends AbstractRoutingDataSource { // [1]

    private Map<Integer, MhaDataSource> shards; // [2]

    @Override
    public void setTargetDataSources(Map<Object, Object> targetDataSources) { // [3]
        super.setTargetDataSources(targetDataSources);

        shards = new HashMap<>();

        for (Object item : targetDataSources.keySet()) {
            String dataSourceName = item.toString();
            String shardNoStr = dataSourceName.split(SHARD_DELIMITER)[0]; // [4]

            MhaDataSource shard = getShard(shardNoStr); // [5]
            if (dataSourceName.contains(MASTER)) {
                shard.setMasterName(dataSourceName);
            } else if (dataSourceName.contains(SLAVE)) {
                shard.getSlaveName().add(dataSourceName);
            }
        }
    }

    @Override
    protected Object determineCurrentLookupKey() { // [6]
        int shardNo = getShardNo(UserHolder.getSharding()); // [7]
        MhaDataSource dataSource = shards.get(shardNo);

        ...

        return TransactionSynchronizationManager.isCurrentTransactionReadOnly() ? dataSource.getSlaveName().next() : dataSource.getMasterName(); // [8]
    }

    private MhaDataSource getShard(String shardNoStr) {
        int shardNo = 0;
        if (StringUtils.isNumeric(shardNoStr)) {
            shardNo = Integer.valueOf(shardNoStr);
        }

        MhaDataSource shard = shards.get(shardNo);
        if (shard == null) {
            shard = new MhaDataSource();
            shard.setSlaveName(new RoundRobin<>(new ArrayList<>()));
            shards.put(shardNo, shard);
        }

        return shard;
    }

    private int getShardNo(UserHolder.Sharding sharding) { // [9]
        if (sharding == null) {
            return 0;
        }

        int shardNo = 0;
        ShardingProperty shardingProperty = ShardingConfig.getShardingPropertyMap().get(sharding.getTarget());
        if (shardingProperty.getStrategy().isRange()) {
            shardNo = getShardNoByRange(shardingProperty.getRules(), sharding.getShardKey());
        } else if (shardingProperty.getStrategy().isModulus()) {
            shardNo = getShardNoByModular(shardingProperty.getMod(), sharding.getShardKey());
        }

        return shardNo;
    }

    private int getShardNoByRange(List<ShardingProperty.ShardingRule> rules, long shardKey) { // [10]
        for (ShardingProperty.ShardingRule rule : rules) {
            if (rule.getRangeMin() <= shardKey && shardKey <= rule.getRangeMax()) {
                return rule.getShardNo();
            }
        }

        return 0;
    }

    private int getShardNoByModular(int modulus, long shardKey) { // [11]
        return (int) (shardKey % modulus);
    }

    @Setter
    @Getter
    private class MhaDataSource { // [12]
        private String masterName;
        private RoundRobin<String> slaveName;
    }
}

public class RoundRobin<T> {
    private final List<T> list;
    private final Iterator<T> iterator;
    private int index;

    public RoundRobin(List<T> list) {
        this.list = list;
        index = 0;
        this.iterator = new Iterator<>() {
            @Override
            public boolean hasNext() {
                return true;
            }

            @Override
            public T next() {
                T value = list.get(index);
                index = (index + 1) % list.size();

                return value;
            }
        };
    }

    public T next() {
        return iterator.next();
    }

    public void add(T item) {
        list.add(item);
    }
}

@Retention(RetentionPolicy.RUNTIME)
@Target(ElementType.TYPE)
public @interface Sharding {
    ShardingTarget target(); // [1]
}

@Service
@CacheConfig(cacheManager = "friendCacheManager")
@Transactional(value = "friendTransactionManager", isolation = Isolation.READ_COMMITTED, propagation = Propagation.REQUIRES_NEW, readOnly = true) // [1]
@Sharding(target = ShardingTarget.FRIEND) // [2]
@RequiredArgsConstructor
public class FriendRepositoryService implements ApplicationContextAware {
    ...

    private final FriendRepository friendRepository;

    @Transactional(value = "friendTransactionManager", isolation = Isolation.READ_COMMITTED) // [3]
    public void save(long userId, FriendEntity entity) { // [4]
        friendRepository.save(entity);
        ...
    }

    @Cacheable(value = FRIEND, key = "{#userId, #friendUserId}", unless = "#result == null")
    public Optional<FriendEntity> findFriend(long userId, long friendUserId) { // [5]
        return friendRepository.findByUserIdAndFriendUserId(userId, friendUserId);
    }

    ...
}

@Component
@Aspect
@RequiredArgsConstructor
public class RepositoryServiceAspect {

    @Pointcut("execution(public * com.woowahan.thiiing.core.repository.service..*.*(..))") // [1]
    private void repositoryService() {
    }

    @Around("repositoryService() && @within(sharding) && args(shardKey,..)") // [2]
    public Object handler(ProceedingJoinPoint pjp, Sharding sharding, long shardKey) throws Throwable {
        UserHolder.setSharding(sharding.target(), shardKey); // [3]

        Object returnVal = pjp.proceed();

        UserHolder.clearSharding(); // [4]

        return returnVal;
    }
}