<commit_before>package org.axonframework.commandhandling.distributed.spring;

import org.axonframework.commandhandling.CommandMessage;
import org.axonframework.commandhandling.distributed.*;
import org.axonframework.commandhandling.distributed.commandfilter.CommandNameFilter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.discovery.DiscoveryClient;
import org.springframework.cloud.client.discovery.event.HeartbeatEvent;
import org.springframework.context.event.EventListener;

import java.net.URI;
import java.util.*;
import java.util.function.Predicate;
import java.util.stream.Collectors;

public class SpringCloudCommandRouter implements CommandRouter {

    private static final Logger LOGGER = LoggerFactory.getLogger(SpringCloudCommandRouter.class);

    private static final String LOAD_FACTOR = "loadFactor";
    private static final String COMMANDS = "commands";
    private static final String COMMA = ",";

    private DiscoveryClient discoveryClient;
    private RoutingStrategy routingStrategy;
    private Set<ServiceInstance> serviceInstances;
    private ConsistentHash consistentHash;


    public SpringCloudCommandRouter(DiscoveryClient discoveryClient, RoutingStrategy routingStrategy) {
        this.discoveryClient = discoveryClient;
        this.routingStrategy = routingStrategy;
        this.serviceInstances = new HashSet<>();
        this.consistentHash = new ConsistentHash();
    }

    /* TODO
    Now recreates the ConsistentHash for every command
    Should consider some mechanism which checks whether members have changed there state

    Additionally, the update done here also takes the localServiceInstance into account, which is currently updated on the updateMembership.
    Hence might lose the `updateMembership(...)` update in favour of the update done in `findDestination(...)`
     */
    @Override
    public Optional<Member> findDestination(CommandMessage<?> commandMessage) {
        LOGGER.info(String.format("findDestination for commandMessage [%s]", commandMessage));

        updateConsistentHash();

        return consistentHash.getMember(routingStrategy.getRoutingKey(commandMessage), commandMessage);
    }

    // Update ConsistentHash for possible updates in supported commands or new nodes
    private void updateConsistentHash() {
        Set<ServiceInstance> currentServiceInstances = discoveryClient.getServices().stream()
                .map(service -> discoveryClient.getInstances(service))
                .flatMap(Collection::stream)
                .collect(Collectors.toSet());

        Set<ServiceInstance> storedServiceInstances = this.serviceInstances;

        currentServiceInstances.removeAll(storedServiceInstances);


        currentServiceInstances.forEach(this::updateConsistentHashWith);
    }

    @EventListener
    public void updateMemberships(HeartbeatEvent event) {
        System.out.println("TODO: Update consistentHash");
    }

    private void updateConsistentHashWith(ServiceInstance serviceInstance) {
        SimpleMember<URI> simpleMember = new SimpleMember<>(serviceInstance.getServiceId(), serviceInstance.getUri(), null);
        Map<String, String> serviceInstanceMetadata = serviceInstance.getMetadata();
        CommandNameFilter commandNameFilter = buildCommandNameFilter(serviceInstanceMetadata.get(COMMANDS));

        consistentHash = consistentHash.with(simpleMember, Integer.parseInt(serviceInstanceMetadata.get(LOAD_FACTOR)), commandNameFilter);
    }

    private CommandNameFilter buildCommandNameFilter(String commaSeparatedCommandsString) {
        return new CommandNameFilter(Arrays.stream(commaSeparatedCommandsString.split(COMMA)).collect(Collectors.toSet()));
    }

    @Override
    public void updateMembership(int loadFactor, Predicate<CommandMessage<?>> commandFilter) {
        LOGGER.info(String.format("updateMembership with loadFactor [%d] and commandFilter predicate [%s]", loadFactor, commandFilter));

        ServiceInstance localServiceInstance = this.discoveryClient.getLocalServiceInstance();
        Map<String, String> localServiceInstanceMetadata = localServiceInstance.getMetadata();
        localServiceInstanceMetadata.put(LOAD_FACTOR, Integer.toString(loadFactor));
        localServiceInstanceMetadata.put(COMMANDS, awkwardInfoStringRemoval(commandFilter.toString()));
    }

    /* TODO
    Preferably we wouldn't have to check for the correct substring to grab, but because the updateMembership function
    in the CommandRouter takes Predicates rather than command names, we have no other way to store them in the
    DiscoveryClient metadata
     */
    private String awkwardInfoStringRemoval(String notReallyCommandSeparatedCommandList) {
        String prependedCommandList = "CommandNameFilter{commandNames=[";
        String appendedToCommandList = "]}";

        int startCommandList = notReallyCommandSeparatedCommandList.indexOf(prependedCommandList) + prependedCommandList.length();
        int endCommandList = notReallyCommandSeparatedCommandList.indexOf(appendedToCommandList);

        return notReallyCommandSeparatedCommandList.substring(startCommandList, endCommandList);
    }

}<commit_msg>Spring Cloud Command Router implementation<commit_after>package org.axonframework.commandhandling.distributed.spring;

import org.axonframework.commandhandling.CommandMessage;
import org.axonframework.commandhandling.distributed.*;
import org.axonframework.commandhandling.distributed.commandfilter.CommandNameFilter;
import org.axonframework.serialization.SerializedObject;
import org.axonframework.serialization.Serializer;
import org.axonframework.serialization.SimpleSerializedObject;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.cloud.client.ServiceInstance;
import org.springframework.cloud.client.discovery.DiscoveryClient;
import org.springframework.cloud.client.discovery.event.HeartbeatEvent;
import org.springframework.context.event.EventListener;

import java.net.URI;
import java.util.*;
import java.util.function.Predicate;
import java.util.stream.Collectors;

public class SpringCloudCommandRouter implements CommandRouter {

    private static final Logger LOGGER = LoggerFactory.getLogger(SpringCloudCommandRouter.class);

    private static final String LOAD_FACTOR = "loadFactor";
    private static final String SERIALIZED_COMMAND_FILTER = "serializedCommandFilter";
    private static final String SERIALIZED_COMMAND_FILTER_CLASS_NAME = "serializedCommandFilterClassName";

    private final DiscoveryClient discoveryClient;
    private final RoutingStrategy routingStrategy;
    private final Serializer serializer;
    private ConsistentHash consistentHash;


    public SpringCloudCommandRouter(DiscoveryClient discoveryClient, RoutingStrategy routingStrategy, Serializer serializer) {
        this.discoveryClient = discoveryClient;
        this.routingStrategy = routingStrategy;
        this.serializer = serializer;
        this.consistentHash = new ConsistentHash();
    }

    @Override
    public Optional<Member> findDestination(CommandMessage<?> commandMessage) {
        LOGGER.info(String.format("findDestination for commandMessage [%s]", commandMessage));

        return consistentHash.getMember(routingStrategy.getRoutingKey(commandMessage), commandMessage);
    }

    @Override
    public void updateMembership(int loadFactor, Predicate<CommandMessage<?>> commandFilter) {
        LOGGER.info(String.format("updateMembership with loadFactor [%d] and commandFilter predicate [%s]", loadFactor, commandFilter));

        ServiceInstance localServiceInstance = this.discoveryClient.getLocalServiceInstance();
        Map<String, String> localServiceInstanceMetadata = localServiceInstance.getMetadata();
        localServiceInstanceMetadata.put(LOAD_FACTOR, Integer.toString(loadFactor));
        SerializedObject<String> serializedCommandFilter = serializer.serialize(commandFilter, String.class);
        localServiceInstanceMetadata.put(SERIALIZED_COMMAND_FILTER, serializedCommandFilter.getData());
        localServiceInstanceMetadata.put(SERIALIZED_COMMAND_FILTER_CLASS_NAME, serializedCommandFilter.getType().getName());

        updateMemberships(Collections.singleton(localServiceInstance));
    }

    @EventListener
    public void updateMemberships(HeartbeatEvent event) {
        LOGGER.info(String.format("update ConsistentHash based on updateMemberships(...) for event [%s]", event));
        Set<ServiceInstance> allServiceInstances = discoveryClient.getServices().stream()
                .map(discoveryClient::getInstances)
                .flatMap(Collection::stream)
                .collect(Collectors.toSet());
        updateMemberships(allServiceInstances);
    }

    private void updateMemberships(Set<ServiceInstance> serviceInstances) {
        Set<Member> currentMembers = consistentHash.getMembers();
        serviceInstances.stream()
                .filter(serviceInstance -> !currentMembers.stream()
                        .filter(memberAlreadyRegistered(serviceInstance))
                        .findFirst()
                        .isPresent()
                ).forEach(this::updateConsistentHashWith);
    }

    // TODO implement a predicate to match on name, URI and command filter
    private Predicate<Member> memberAlreadyRegistered(ServiceInstance serviceInstance) {
//        return member -> serviceInstance.getServiceId().equals(member.name());
        return member -> false;
    }

    private void updateConsistentHashWith(ServiceInstance newServiceInstance) {
        SimpleMember<URI> simpleMember = new SimpleMember<>(newServiceInstance.getServiceId(),
                newServiceInstance.getUri(), member -> consistentHash.without(member));

        Map<String, String> serviceInstanceMetadata = newServiceInstance.getMetadata();

        int loadFactor = Integer.parseInt(serviceInstanceMetadata.get(LOAD_FACTOR));
        SimpleSerializedObject<String> serializedObject = new SimpleSerializedObject<>(
                serviceInstanceMetadata.get(SERIALIZED_COMMAND_FILTER), String.class,
                serviceInstanceMetadata.get(SERIALIZED_COMMAND_FILTER_CLASS_NAME), null);
        CommandNameFilter commandNameFilter = serializer.deserialize(serializedObject);

        consistentHash = consistentHash.with(simpleMember, loadFactor, commandNameFilter);
    }

}
{'fraction_non_alphanumeric': 0.0799388028303691, 'fraction_numerical': 0.00038248231019315357, 'mean_word_length': 4.678067318132465, 'pattern_counts': {'":': 0, '<': 30, '<?xml version=': 0, '>': 35, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3639949', 'n_tokens_mistral': 2546, 'n_tokens_neox': 2444, 'n_words': 533}