diff --git a/main/query-support/src/main/java/org/elasticsoftware/akces/query/reflector/AkcesReflectorController.java b/main/query-support/src/main/java/org/elasticsoftware/akces/query/reflector/AkcesReflectorController.java new file mode 100644 index 00000000..9c600c3a --- /dev/null +++ b/main/query-support/src/main/java/org/elasticsoftware/akces/query/reflector/AkcesReflectorController.java @@ -0,0 +1,424 @@ +/* + * Copyright 2022 - 2026 The Original 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.elasticsoftware.akces.query.reflector; + +import tools.jackson.databind.ObjectMapper; +import com.google.common.hash.HashFunction; +import com.google.common.hash.Hashing; +import jakarta.annotation.Nonnull; +import org.apache.kafka.clients.consumer.Consumer; +import org.apache.kafka.clients.consumer.ConsumerRebalanceListener; +import org.apache.kafka.clients.consumer.ConsumerRecords; +import org.apache.kafka.common.KafkaException; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.errors.InterruptException; +import org.apache.kafka.common.errors.WakeupException; +import org.elasticsoftware.akces.aggregate.CommandType; +import org.elasticsoftware.akces.aggregate.DomainEventType; +import org.elasticsoftware.akces.annotations.CommandInfo; +import org.elasticsoftware.akces.commands.Command; +import org.elasticsoftware.akces.control.AggregateServiceCommandType; +import org.elasticsoftware.akces.control.AggregateServiceDomainEventType; +import org.elasticsoftware.akces.control.AggregateServiceRecord; +import org.elasticsoftware.akces.control.AkcesControlRecord; +import org.elasticsoftware.akces.control.AkcesRegistry; +import org.elasticsoftware.akces.gdpr.GDPRContextRepositoryFactory; +import org.elasticsoftware.akces.protocol.ProtocolRecord; +import org.elasticsoftware.akces.protocol.SchemaRecord; +import org.elasticsoftware.akces.schemas.KafkaSchemaRegistry; +import org.elasticsoftware.akces.schemas.SchemaRegistry; +import org.elasticsoftware.akces.schemas.storage.KafkaTopicSchemaStorage; +import org.elasticsoftware.akces.schemas.storage.SchemaStorage; +import org.elasticsoftware.akces.util.HostUtils; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.beans.BeansException; +import org.springframework.boot.availability.AvailabilityChangeEvent; +import org.springframework.boot.availability.LivenessState; +import org.springframework.context.ApplicationContext; +import org.springframework.context.ApplicationContextAware; +import org.springframework.kafka.core.ConsumerFactory; +import org.springframework.kafka.core.ProducerFactory; +import org.springframework.scheduling.concurrent.CustomizableThreadFactory; + +import java.time.Duration; +import java.util.*; +import java.util.concurrent.*; + +import static java.nio.charset.StandardCharsets.UTF_8; +import static org.elasticsoftware.akces.gdpr.GDPRAnnotationUtils.hasPIIDataAnnotation; +import static org.elasticsoftware.akces.query.reflector.AkcesReflectorControllerState.*; + +public class AkcesReflectorController extends Thread + implements AutoCloseable, ConsumerRebalanceListener, ApplicationContextAware, AkcesRegistry { + + private static final Logger logger = LoggerFactory.getLogger(AkcesReflectorController.class); + + private final ConsumerFactory consumerFactory; + private final ProducerFactory producerFactory; + private final ConsumerFactory controlRecordConsumerFactory; + private final ConsumerFactory schemaRecordConsumerFactory; + private final ObjectMapper objectMapper; + private final GDPRContextRepositoryFactory gdprContextRepositoryFactory; + private final ReflectorRuntime reflectorRuntime; + private final Map reflectorPartitions = new HashMap<>(); + private final ExecutorService executorService; + private final HashFunction hashFunction = Hashing.murmur3_32_fixed(); + private final Map aggregateServices = new ConcurrentHashMap<>(); + private final List partitionsToAssign = new ArrayList<>(); + private final List partitionsToRevoke = new ArrayList<>(); + private final CountDownLatch shutdownLatch = new CountDownLatch(1); + private Integer partitions = null; + private Consumer controlConsumer; + private volatile AkcesReflectorControllerState processState = INITIALIZING; + private ApplicationContext applicationContext; + private SchemaStorage schemaStorage; + private SchemaRegistry schemaRegistry; + + public AkcesReflectorController(ConsumerFactory consumerFactory, + ProducerFactory producerFactory, + ConsumerFactory controlConsumerFactory, + ConsumerFactory schemaRecordConsumerFactory, + ObjectMapper objectMapper, + GDPRContextRepositoryFactory gdprContextRepositoryFactory, + ReflectorRuntime reflectorRuntime) { + super(reflectorRuntime.getName() + "-AkcesReflectorController"); + this.consumerFactory = consumerFactory; + this.producerFactory = producerFactory; + this.controlRecordConsumerFactory = controlConsumerFactory; + this.schemaRecordConsumerFactory = schemaRecordConsumerFactory; + this.objectMapper = objectMapper; + this.gdprContextRepositoryFactory = gdprContextRepositoryFactory; + this.reflectorRuntime = reflectorRuntime; + this.executorService = Executors.newCachedThreadPool( + new CustomizableThreadFactory(reflectorRuntime.getName() + "ReflectorPartitionThread-")); + } + + @Override + public void run() { + try { + controlConsumer = controlRecordConsumerFactory.createConsumer( + reflectorRuntime.getName() + "Reflector-Akces-Control", + reflectorRuntime.getName() + "-" + HostUtils.getHostName() + "Reflector-Akces-Control", + null); + controlConsumer.subscribe(List.of("Akces-Control"), this); + // determine the number of partitions once metadata is available + this.partitions = controlConsumer.partitionsFor("Akces-Control").size(); + this.schemaStorage = new KafkaTopicSchemaStorage( + schemaRecordConsumerFactory.createConsumer( + HostUtils.getHostName() + "-AkcesReflectorController-Schemas", + HostUtils.getHostName() + "-AkcesReflectorController-Schemas", + null)); + schemaStorage.initialize(); + this.schemaRegistry = new KafkaSchemaRegistry(schemaStorage); + // validate domain event schemas + reflectorRuntime.validateDomainEventSchemas(schemaRegistry); + while (processState != SHUTTING_DOWN) { + process(); + } + // close all reflector partitions + logger.info("Closing {} ReflectorPartitions", reflectorPartitions.size()); + reflectorPartitions.values().forEach(reflectorPartition -> { + if (reflectorPartition != null) { + try { + reflectorPartition.close(); + } catch (Exception e) { + logger.error("Error closing ReflectorPartition " + reflectorPartition.getId(), e); + } + } + }); + try { + controlConsumer.close(Duration.ofSeconds(5)); + } catch (InterruptException e) { + Thread.currentThread().interrupt(); + } catch (KafkaException e) { + logger.error("Error closing controlConsumer", e); + } + // raise an error for the liveness check + applicationContext.publishEvent(new AvailabilityChangeEvent<>(this, LivenessState.BROKEN)); + // signal done + shutdownLatch.countDown(); + } catch (Exception e) { + logger.error("Error in AkcesReflectorController", e); + processState = ERROR; + } + } + + private void process() { + if (processState == RUNNING) { + try { + processControlRecords(); + } catch (WakeupException e) { + // ignore + } catch (InterruptException e) { + Thread.currentThread().interrupt(); + } catch (KafkaException e) { + logger.error("Unrecoverable exception in AkcesReflectorController", e); + processState = SHUTTING_DOWN; + } + } else if (processState == INITIALIZING) { + processControlRecords(); + // we don't switch to the running state here; we wait for the INITIAL_REBALANCING state + // which should be triggered by onPartitionsRevoked/onPartitionsAssigned + } else if (processState == INITIAL_REBALANCING) { + // load all the service data from the control topic before transitioning to REBALANCING + if (!partitionsToAssign.isEmpty()) { + try { + controlConsumer.seekToBeginning(partitionsToAssign); + Map initializedEndOffsets = controlConsumer.endOffsets(partitionsToAssign); + ConsumerRecords consumerRecords = + controlConsumer.poll(Duration.ofMillis(100)); + while (!initializedEndOffsets.isEmpty()) { + consumerRecords.forEach(record -> { + AkcesControlRecord controlRecord = record.value(); + if (controlRecord instanceof AggregateServiceRecord aggregateServiceRecord) { + if (aggregateServices.putIfAbsent(record.key(), aggregateServiceRecord) == null) { + logger.info("Discovered service: {}", aggregateServiceRecord.aggregateName()); + } + } else { + logger.info("Received unknown AkcesControlRecord type: {}", + controlRecord.getClass().getSimpleName()); + } + }); + if (consumerRecords.isEmpty()) { + initializedEndOffsets.entrySet().removeIf( + entry -> entry.getValue() <= controlConsumer.position(entry.getKey())); + } + consumerRecords = controlConsumer.poll(Duration.ofMillis(100)); + } + } catch (WakeupException e) { + // ignore + } catch (InterruptException e) { + Thread.currentThread().interrupt(); + } catch (KafkaException e) { + logger.error("Unrecoverable exception in AkcesReflectorController", e); + processState = SHUTTING_DOWN; + } + } + // now we can move to REBALANCING + processState = REBALANCING; + } else if (processState == REBALANCING) { + // first revoke + for (TopicPartition topicPartition : partitionsToRevoke) { + ReflectorPartition reflectorPartition = reflectorPartitions.remove(topicPartition.partition()); + if (reflectorPartition != null) { + logger.info("Stopping ReflectorPartition {}", reflectorPartition.getId()); + try { + reflectorPartition.close(); + } catch (Exception e) { + logger.error("Error closing ReflectorPartition", e); + } + } + } + partitionsToRevoke.clear(); + // then assign + for (TopicPartition topicPartition : partitionsToAssign) { + ReflectorPartition reflectorPartition = new ReflectorPartition( + consumerFactory, + producerFactory, + reflectorRuntime, + gdprContextRepositoryFactory, + objectMapper, + topicPartition.partition(), + new TopicPartition("Akces-GDPRKeys", topicPartition.partition()), + reflectorRuntime.getDomainEventTypes(), + this); + reflectorPartitions.put(reflectorPartition.getId(), reflectorPartition); + logger.info("Starting ReflectorPartition {}", reflectorPartition.getId()); + executorService.submit(reflectorPartition); + } + partitionsToAssign.clear(); + // move back to running + processState = RUNNING; + } + } + + private void processControlRecords() { + ConsumerRecords consumerRecords = + controlConsumer.poll(Duration.ofMillis(100)); + if (!consumerRecords.isEmpty()) { + consumerRecords.forEach(record -> { + AkcesControlRecord controlRecord = record.value(); + if (controlRecord instanceof AggregateServiceRecord aggregateServiceRecord) { + if (!aggregateServices.containsKey(record.key())) { + logger.info("Discovered service: {}", aggregateServiceRecord.aggregateName()); + } + // always overwrite with the latest version + aggregateServices.put(record.key(), aggregateServiceRecord); + } else { + logger.info("Received unknown AkcesControlRecord type: {}", + controlRecord.getClass().getSimpleName()); + } + }); + } + } + + @Override + public void close() { + logger.info("Shutting down AkcesReflectorController"); + this.processState = SHUTTING_DOWN; + try { + if (shutdownLatch.await(10, TimeUnit.SECONDS)) { + logger.info("AkcesReflectorController has been shutdown"); + } else { + logger.warn("AkcesReflectorController did not shutdown within 10 seconds"); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + + @Override + public void onPartitionsRevoked(Collection topicPartitions) { + if (!topicPartitions.isEmpty()) { + partitionsToRevoke.addAll(topicPartitions); + if (processState == RUNNING) { + logger.info("Switching from RUNNING to REBALANCING, revoking partitions: {}", + topicPartitions.stream().map(TopicPartition::partition).toList()); + processState = REBALANCING; + } else if (processState == INITIALIZING) { + logger.info("Switching from INITIALIZING to INITIAL_REBALANCING, revoking partitions: {}", + topicPartitions.stream().map(TopicPartition::partition).toList()); + processState = INITIAL_REBALANCING; + } + } + } + + @Override + public void onPartitionsAssigned(Collection topicPartitions) { + if (!topicPartitions.isEmpty()) { + partitionsToAssign.addAll(topicPartitions); + if (processState == RUNNING) { + logger.info("Switching from RUNNING to REBALANCING, assigning partitions: {}", + topicPartitions.stream().map(TopicPartition::partition).toList()); + processState = REBALANCING; + } else if (processState == INITIALIZING) { + logger.info("Switching from INITIALIZING to INITIAL_REBALANCING, assigning partitions: {}", + topicPartitions.stream().map(TopicPartition::partition).toList()); + processState = INITIAL_REBALANCING; + } + } + } + + private boolean supportsCommand(List supportedCommands, CommandInfo commandInfo) { + for (AggregateServiceCommandType supportedCommand : supportedCommands) { + if (supportedCommand.typeName().equals(commandInfo.type()) && + supportedCommand.version() == commandInfo.version()) { + return true; + } + } + return false; + } + + private boolean supportsCommand(List supportedCommands, CommandType commandType) { + for (AggregateServiceCommandType supportedCommand : supportedCommands) { + if (supportedCommand.typeName().equals(commandType.typeName()) && + supportedCommand.version() == commandType.version()) { + return true; + } + } + return false; + } + + private boolean producesDomainEvent(List producedEvents, + DomainEventType externalDomainEventType) { + for (AggregateServiceDomainEventType producedEvent : producedEvents) { + if (producedEvent.typeName().equals(externalDomainEventType.typeName()) && + producedEvent.version() == externalDomainEventType.version()) { + return true; + } + } + return false; + } + + @Override + @Nonnull + public CommandType resolveType(@Nonnull Class commandClass) { + CommandInfo commandInfo = commandClass.getAnnotation(CommandInfo.class); + if (commandInfo == null) { + throw new IllegalStateException( + "Command class " + commandClass.getName() + " is not annotated with @CommandInfo"); + } + List services = aggregateServices.values().stream() + .filter(record -> supportsCommand(record.supportedCommands(), commandInfo)) + .toList(); + if (services.size() == 1) { + return new CommandType<>( + commandInfo.type(), + commandInfo.version(), + commandClass, + false, + true, + hasPIIDataAnnotation(commandClass)); + } else { + throw new IllegalStateException( + "Cannot determine where to send command " + commandClass.getName()); + } + } + + @Override + @Nonnull + public String resolveTopic(@Nonnull Class commandClass) { + return resolveTopic(resolveType(commandClass)); + } + + @Override + @Nonnull + public String resolveTopic(@Nonnull CommandType commandType) { + List services = aggregateServices.values().stream() + .filter(record -> supportsCommand(record.supportedCommands(), commandType)) + .toList(); + if (services.size() == 1) { + return services.getFirst().commandTopic(); + } else { + throw new IllegalStateException( + "Cannot determine where to send command " + commandType.typeName() + " v" + commandType.version()); + } + } + + @Override + @Nonnull + public String resolveTopic(@Nonnull DomainEventType externalDomainEventType) { + List services = aggregateServices.values().stream() + .filter(record -> producesDomainEvent(record.producedEvents(), externalDomainEventType)) + .toList(); + if (services.size() == 1) { + return services.getFirst().domainEventTopic(); + } else { + throw new IllegalStateException( + "Cannot determine which service produces DomainEvent " + + externalDomainEventType.typeName() + " v" + externalDomainEventType.version()); + } + } + + @Override + @Nonnull + public Integer resolvePartition(@Nonnull String aggregateId) { + return Math.abs(hashFunction.hashString(aggregateId, UTF_8).asInt()) % partitions; + } + + public boolean isRunning() { + return processState == RUNNING + && reflectorPartitions.values().stream().allMatch(ReflectorPartition::isProcessing); + } + + @Override + public void setApplicationContext(ApplicationContext applicationContext) throws BeansException { + this.applicationContext = applicationContext; + } +} diff --git a/main/query-support/src/main/java/org/elasticsoftware/akces/query/reflector/AkcesReflectorControllerState.java b/main/query-support/src/main/java/org/elasticsoftware/akces/query/reflector/AkcesReflectorControllerState.java new file mode 100644 index 00000000..4bafee54 --- /dev/null +++ b/main/query-support/src/main/java/org/elasticsoftware/akces/query/reflector/AkcesReflectorControllerState.java @@ -0,0 +1,27 @@ +/* + * Copyright 2022 - 2026 The Original 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.elasticsoftware.akces.query.reflector; + +public enum AkcesReflectorControllerState { + INITIALIZING, + INITIAL_REBALANCING, + REBALANCING, + RUNNING, + SHUTTING_DOWN, + ERROR +} diff --git a/main/query-support/src/main/java/org/elasticsoftware/akces/query/reflector/ReflectorPartition.java b/main/query-support/src/main/java/org/elasticsoftware/akces/query/reflector/ReflectorPartition.java new file mode 100644 index 00000000..c1e3ba5d --- /dev/null +++ b/main/query-support/src/main/java/org/elasticsoftware/akces/query/reflector/ReflectorPartition.java @@ -0,0 +1,396 @@ +/* + * Copyright 2022 - 2026 The Original 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.elasticsoftware.akces.query.reflector; + +import tools.jackson.databind.ObjectMapper; +import org.apache.kafka.clients.consumer.Consumer; +import org.apache.kafka.clients.consumer.ConsumerRecord; +import org.apache.kafka.clients.consumer.ConsumerRecords; +import org.apache.kafka.clients.consumer.OffsetAndMetadata; +import org.apache.kafka.clients.producer.Producer; +import org.apache.kafka.clients.producer.ProducerRecord; +import org.apache.kafka.common.KafkaException; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.errors.InterruptException; +import org.apache.kafka.common.errors.InvalidProducerEpochException; +import org.apache.kafka.common.errors.WakeupException; +import org.elasticsoftware.akces.aggregate.CommandType; +import org.elasticsoftware.akces.aggregate.DomainEventType; +import org.elasticsoftware.akces.commands.Command; +import org.elasticsoftware.akces.commands.CommandBus; +import org.elasticsoftware.akces.control.AkcesRegistry; +import org.elasticsoftware.akces.gdpr.GDPRContextRepository; +import org.elasticsoftware.akces.gdpr.GDPRContextRepositoryFactory; +import org.elasticsoftware.akces.protocol.CommandRecord; +import org.elasticsoftware.akces.protocol.DomainEventRecord; +import org.elasticsoftware.akces.protocol.PayloadEncoding; +import org.elasticsoftware.akces.protocol.ProtocolRecord; +import org.elasticsoftware.akces.util.HostUtils; +import org.elasticsoftware.akces.util.KafkaSender; +import org.slf4j.Logger; +import org.slf4j.LoggerFactory; +import org.springframework.kafka.core.ConsumerFactory; +import org.springframework.kafka.core.ProducerFactory; + +import java.io.IOException; +import java.time.Duration; +import java.time.Instant; +import java.util.*; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.stream.Collectors; +import java.util.stream.Stream; + +import static java.util.Collections.singletonList; +import static org.elasticsoftware.akces.query.reflector.ReflectorPartitionState.*; +import static org.elasticsoftware.akces.query.reflector.ReflectorResult.*; + +public class ReflectorPartition implements Runnable, AutoCloseable, CommandBus { + private static final Logger logger = LoggerFactory.getLogger(ReflectorPartition.class); + + private final ConsumerFactory consumerFactory; + private final ProducerFactory producerFactory; + private final ReflectorRuntime runtime; + private final GDPRContextRepository gdprContextRepository; + private final ObjectMapper objectMapper; + private final Integer id; + private final Set externalEventPartitions = new HashSet<>(); + private final Collection> externalDomainEventTypes; + private final CountDownLatch shutdownLatch = new CountDownLatch(1); + private final AkcesRegistry akcesRegistry; + private final TopicPartition gdprKeyPartition; + /** Maps paused partitions to the earliest instant at which they may be resumed. */ + private final Map pausedUntil = new HashMap<>(); + /** Tracks the current retry attempt count per partition for linear backoff calculation. */ + private final Map retryAttempts = new HashMap<>(); + + private Consumer consumer; + private Producer producer; + private volatile ReflectorPartitionState processState; + private Map initializedEndOffsets = Collections.emptyMap(); + private volatile Thread reflectorPartitionThread = null; + + public ReflectorPartition(ConsumerFactory consumerFactory, + ProducerFactory producerFactory, + ReflectorRuntime runtime, + GDPRContextRepositoryFactory gdprContextRepositoryFactory, + ObjectMapper objectMapper, + Integer id, + TopicPartition gdprKeyPartition, + Collection> externalDomainEventTypes, + AkcesRegistry akcesRegistry) { + this.consumerFactory = consumerFactory; + this.producerFactory = producerFactory; + this.runtime = runtime; + this.objectMapper = objectMapper; + this.id = id; + this.externalDomainEventTypes = externalDomainEventTypes; + this.akcesRegistry = akcesRegistry; + this.gdprKeyPartition = gdprKeyPartition; + this.processState = INITIALIZING; + this.gdprContextRepository = gdprContextRepositoryFactory.create(runtime.getName(), id); + } + + public Integer getId() { + return id; + } + + @Override + public void run() { + try { + this.reflectorPartitionThread = Thread.currentThread(); + ReflectorPartitionCommandBus.registerCommandBus(this); + logger.info("Starting ReflectorPartition {} of {}Reflector", id, runtime.getName()); + this.consumer = consumerFactory.createConsumer( + runtime.getName() + "Reflector-partition-" + id, + runtime.getName() + "Reflector-partition-" + id + "-" + HostUtils.getHostName(), + null); + this.producer = producerFactory.createProducer( + runtime.getName() + "Reflector-partition-" + id + "-" + HostUtils.getHostName()); + // resolve the external event partitions + Set distinctTopics = externalDomainEventTypes.stream() + .map(akcesRegistry::resolveTopic) + .collect(Collectors.toSet()); + externalEventPartitions.addAll(distinctTopics.stream() + .map(topic -> new TopicPartition(topic, id)) + .collect(Collectors.toSet())); + // make a hard assignment, only assign the gdprKeyPartition if needed + consumer.assign(Stream.concat( + runtime.shouldHandlePIIData() ? Stream.of(gdprKeyPartition) : Stream.empty(), + externalEventPartitions.stream()).toList()); + logger.info("Assigned partitions {} for ReflectorPartition {} of {}Reflector", + consumer.assignment(), id, runtime.getName()); + while (processState != SHUTTING_DOWN) { + process(); + } + logger.info("Shutting down ReflectorPartition {} of {}Reflector", id, runtime.getName()); + } catch (Throwable t) { + logger.error("Unexpected error in ReflectorPartition {} of {}Reflector", id, runtime.getName(), t); + } finally { + try { + consumer.close(Duration.ofSeconds(5)); + producer.close(Duration.ofSeconds(5)); + } catch (InterruptException e) { + Thread.currentThread().interrupt(); + } catch (KafkaException e) { + logger.error("Error closing consumer/producer", e); + } + try { + gdprContextRepository.close(); + } catch (IOException e) { + logger.error("Error closing gdpr context repository", e); + } + ReflectorPartitionCommandBus.registerCommandBus(null); + } + logger.info("Finished Shutting down ReflectorPartition {} of {}Reflector", id, runtime.getName()); + shutdownLatch.countDown(); + } + + @Override + public void close() { + processState = SHUTTING_DOWN; + // wait maximum of 10 seconds for the shutdown to complete + try { + if (shutdownLatch.await(10, TimeUnit.SECONDS)) { + logger.info("ReflectorPartition={} has been shutdown", id); + } else { + logger.warn("ReflectorPartition={} did not shutdown within 10 seconds", id); + } + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + } + } + + @Override + public void send(Command command) { + // this implementation is only meant to be called from the ReflectorPartition thread + if (Thread.currentThread() != reflectorPartitionThread) { + throw new IllegalStateException("send() can only be called from the ReflectorPartition thread"); + } + CommandType commandType = akcesRegistry.resolveType(command.getClass()); + if (commandType != null) { + String topic = akcesRegistry.resolveTopic(commandType); + byte[] payload; + try { + payload = objectMapper.writeValueAsBytes(command); + } catch (Exception e) { + logger.error("Error serializing command {}", commandType.typeName(), e); + throw new RuntimeException(e); + } + CommandRecord commandRecord = new CommandRecord( + null, + commandType.typeName(), + commandType.version(), + payload, + PayloadEncoding.JSON, + command.getAggregateId(), + null, + null); // don't send a response + Integer partition = akcesRegistry.resolvePartition(command.getAggregateId()); + KafkaSender.send(producer, new ProducerRecord<>(topic, partition, commandRecord.id(), commandRecord)); + } + } + + private void process() { + try { + if (processState == PROCESSING) { + // resume any paused partitions whose backoff has elapsed + resumeReadyPartitions(); + ConsumerRecords allRecords = consumer.poll(Duration.ofMillis(10)); + if (!allRecords.isEmpty()) { + processRecords(allRecords); + } + } else if (processState == LOADING_GDPR_KEYS) { + ConsumerRecords gdprKeyRecords = consumer.poll(Duration.ofMillis(10)); + gdprContextRepository.process(gdprKeyRecords.records(gdprKeyPartition)); + // stop condition + if (gdprKeyRecords.isEmpty() && initializedEndOffsets.getOrDefault(gdprKeyPartition, 0L) <= consumer.position(gdprKeyPartition)) { + // resume the other topics + consumer.resume(externalEventPartitions); + processState = PROCESSING; + } + } else if (processState == INITIALIZING) { + logger.info( + "Initializing ReflectorPartition {} of {}Reflector. Will {}", + id, + runtime.getName(), + runtime.shouldHandlePIIData() ? "Handle PII Data" : "Not Handle PII Data"); + // handle possible GDPRContext setup + if (runtime.shouldHandlePIIData()) { + long gdprKeyRepositoryOffset = gdprContextRepository.getOffset(); + if (gdprKeyRepositoryOffset >= 0) { + logger.info( + "Resuming GDPRKeys from offset {} for ReflectorPartition {} of {}Reflector", + gdprKeyRepositoryOffset, + id, + runtime.getName()); + consumer.seek(gdprKeyPartition, gdprContextRepository.getOffset() + 1); + } else { + consumer.seekToBeginning(singletonList(gdprKeyPartition)); + } + // find the end offsets so we know when to stop + initializedEndOffsets = consumer.endOffsets(List.of(gdprKeyPartition)); + logger.info("Loading GDPR Keys for ReflectorPartition {} of {}Reflector", id, runtime.getName()); + // pause the external event partitions while loading GDPR keys + consumer.pause(externalEventPartitions); + processState = LOADING_GDPR_KEYS; + } else { + // skip the LOADING_GDPR_KEYS phase and go directly to PROCESSING + consumer.resume(externalEventPartitions); + processState = PROCESSING; + } + } + } catch (WakeupException | InterruptException ignore) { + // non-fatal, ignore + } catch (KafkaException e) { + // fatal + logger.error("Fatal error during {} phase, shutting down ReflectorPartition {} of {}Reflector", + processState, id, runtime.getName(), e); + processState = SHUTTING_DOWN; + } + } + + private void resumeReadyPartitions() { + if (!pausedUntil.isEmpty()) { + Instant now = Instant.now(); + Set toResume = new HashSet<>(); + pausedUntil.entrySet().removeIf(entry -> { + if (now.isAfter(entry.getValue())) { + toResume.add(entry.getKey()); + return true; + } + return false; + }); + if (!toResume.isEmpty()) { + logger.debug("Resuming {} paused partition(s) for ReflectorPartition {} of {}Reflector", + toResume.size(), id, runtime.getName()); + consumer.resume(toResume); + } + } + } + + private void processRecords(ConsumerRecords allRecords) { + if (logger.isTraceEnabled()) { + logger.trace("Processing {} records in poll batch for ReflectorPartition {} of {}Reflector", + allRecords.count(), id, runtime.getName()); + } + // first handle GDPR key records (outside of the event transaction) + if (runtime.shouldHandlePIIData()) { + List> gdprKeyRecords = allRecords.records(gdprKeyPartition); + if (!gdprKeyRecords.isEmpty()) { + gdprContextRepository.process(gdprKeyRecords); + } + } + // process domain events one at a time, each in its own transaction + for (TopicPartition partition : allRecords.partitions()) { + if (partition.equals(gdprKeyPartition)) { + continue; // already handled above + } + for (ConsumerRecord record : allRecords.records(partition)) { + if (record.value() instanceof DomainEventRecord domainEventRecord) { + processEventRecord(domainEventRecord, partition, record.offset()); + } + } + } + } + + private void processEventRecord(DomainEventRecord domainEventRecord, + TopicPartition topicPartition, + long offset) { + ReflectorResult result; + try { + result = runtime.apply(domainEventRecord, topicPartition, gdprContextRepository::get); + } catch (IOException e) { + logger.error("IOException deserializing event '{}' on partition {} of {}Reflector", + domainEventRecord.name(), topicPartition, runtime.getName(), e); + // treat deserialization failure as a permanent failure — skip the event + result = FAILURE; + } + + switch (result) { + case RETRY_PENDING -> { + // backoff has not elapsed yet; pause the partition and seek back so the event is re-delivered + int attempt = retryAttempts.merge(topicPartition, 1, Integer::sum); + long backoffMs = (long) attempt * runtime.getRetryBackoffBaseMs(); + pausedUntil.put(topicPartition, Instant.now().plusMillis(backoffMs)); + consumer.pause(Set.of(topicPartition)); + consumer.seek(topicPartition, offset); + logger.debug("ReflectorPartition {} of {}Reflector: RETRY_PENDING for event '{}' on partition {} " + + "(attempt {}); pausing for {} ms", + id, runtime.getName(), domainEventRecord.name(), topicPartition, attempt, backoffMs); + } + case SUCCESS -> { + retryAttempts.remove(topicPartition); + pausedUntil.remove(topicPartition); + commitEventWithTransaction( + domainEventRecord, topicPartition, offset, + true, runtime.getLastEvent(), runtime.getLastResult(), null); + } + case FAILURE -> { + retryAttempts.remove(topicPartition); + pausedUntil.remove(topicPartition); + commitEventWithTransaction( + domainEventRecord, topicPartition, offset, + false, runtime.getLastEvent(), null, runtime.getLastException()); + } + } + } + + private void commitEventWithTransaction(DomainEventRecord domainEventRecord, + TopicPartition topicPartition, + long offset, + boolean success, + org.elasticsoftware.akces.events.DomainEvent event, + Object result, + Exception exception) { + try { + producer.beginTransaction(); + try { + if (success) { + runtime.handleSuccess(event, result, this); + } else { + runtime.handleFailure(event, exception, this); + } + producer.sendOffsetsToTransaction( + Map.of(topicPartition, new OffsetAndMetadata(offset + 1)), + consumer.groupMetadata()); + producer.commitTransaction(); + gdprContextRepository.commit(); + } catch (InvalidProducerEpochException e) { + producer.abortTransaction(); + consumer.seek(topicPartition, offset); + gdprContextRepository.rollback(); + logger.warn("Transaction aborted for event '{}' on partition {} of {}Reflector due to InvalidProducerEpochException", + domainEventRecord.name(), topicPartition, runtime.getName(), e); + } catch (KafkaException e) { + producer.abortTransaction(); + consumer.seek(topicPartition, offset); + gdprContextRepository.rollback(); + logger.error("Transaction aborted for event '{}' on partition {} of {}Reflector", + domainEventRecord.name(), topicPartition, runtime.getName(), e); + } + } catch (KafkaException e) { + logger.error("Error beginning transaction for event '{}' on partition {} of {}Reflector", + domainEventRecord.name(), topicPartition, runtime.getName(), e); + } + } + + public boolean isProcessing() { + return processState == PROCESSING; + } +} diff --git a/main/query-support/src/main/java/org/elasticsoftware/akces/query/reflector/ReflectorPartitionCommandBus.java b/main/query-support/src/main/java/org/elasticsoftware/akces/query/reflector/ReflectorPartitionCommandBus.java new file mode 100644 index 00000000..b7ddaf43 --- /dev/null +++ b/main/query-support/src/main/java/org/elasticsoftware/akces/query/reflector/ReflectorPartitionCommandBus.java @@ -0,0 +1,26 @@ +/* + * Copyright 2022 - 2026 The Original 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.elasticsoftware.akces.query.reflector; + +import org.elasticsoftware.akces.commands.CommandBusHolder; + +class ReflectorPartitionCommandBus extends CommandBusHolder { + static void registerCommandBus(ReflectorPartition reflectorPartition) { + commandBusThreadLocal.set(reflectorPartition); + } +} diff --git a/main/query-support/src/main/java/org/elasticsoftware/akces/query/reflector/ReflectorPartitionState.java b/main/query-support/src/main/java/org/elasticsoftware/akces/query/reflector/ReflectorPartitionState.java new file mode 100644 index 00000000..864a3743 --- /dev/null +++ b/main/query-support/src/main/java/org/elasticsoftware/akces/query/reflector/ReflectorPartitionState.java @@ -0,0 +1,25 @@ +/* + * Copyright 2022 - 2026 The Original 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.elasticsoftware.akces.query.reflector; + +public enum ReflectorPartitionState { + INITIALIZING, + LOADING_GDPR_KEYS, + PROCESSING, + SHUTTING_DOWN +}