package com.sixt.service.framework.kafka;

import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.common.Metric;
import org.apache.kafka.common.MetricName;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import javax.inject.Singleton;
import java.time.Instant;
import java.time.temporal.ChronoUnit;
import java.util.*;
import java.util.concurrent.CopyOnWriteArraySet;
import java.util.concurrent.Executors;
import java.util.concurrent.ScheduledFuture;
import java.util.concurrent.TimeUnit;

@Singleton
public class PartitionAssignmentWatchdog {

    private static final Logger logger = LoggerFactory.getLogger(PartitionAssignmentWatchdog.class);

    public static final int SECONDS_BEFORE_KILLING = 30;

    private Set<KafkaConsumer<String, String>> consumers = new CopyOnWriteArraySet<>();
    private Map<String, ConsumerPartitionLevels> consumerPartitions = new HashMap<>();
    private ScheduledFuture<?> partitionWatchdogFuture;

    public synchronized void subscriberInitialized(KafkaConsumer<String, String> realConsumer) {
        logger.debug("Adding subscriber");
        consumers.add(realConsumer);
        if (consumers.size() == 1) {
            startWatchdog();
        }
    }

    public synchronized void subscriberShutdown(KafkaConsumer<String, String> realConsumer) {
        logger.debug("Removing subscriber");
        consumers.remove(realConsumer);
        if (consumers.size() == 0) {
            stopWatchdog();
        }
    }

    private synchronized void checkAssignments() {
        logger.debug("Checking partition assignments");

        try {
            KafkaConsumer<String, String> consumer = consumers.iterator().next();
            Map<MetricName, ? extends Metric> metrics = consumer.metrics();
            for (MetricName name : metrics.keySet()) {
                if ("assigned-partitions".equals(name.name())) {
                    Metric metric = metrics.get(name);
                    Map<String, String> tags = name.tags();
                    String clientId = tags.get("client-id");
                    int partitionCount = ((Double)metric.metricValue()).intValue();
                    processDataPoint(clientId, partitionCount, Instant.now());
                }
            }
        } catch (NoSuchElementException ex) {
        }
    }

    private void processDataPoint(String clientId, int partitionCount, Instant now) {
        consumerPartitions.putIfAbsent(clientId, new ConsumerPartitionLevels());
        ConsumerPartitionLevels clientInfo = consumerPartitions.get(clientId);
        if (partitionCount != clientInfo.lastKnownCount) {
            clientInfo.lastKnownCount = partitionCount;
            clientInfo.lastChangeDate = now;
        }
        if (clientInfo.lastKnownCount == 0 && clientInfo.lastChangeDate
                .isBefore(now.minus(SECONDS_BEFORE_KILLING, ChronoUnit.SECONDS))) {
            logger.error("TESTING: would now kill the application because consumer {} has had no " +
                    "partition assignments for more than {} seconds", clientId, SECONDS_BEFORE_KILLING);
        }
    }

    private void startWatchdog() {
        partitionWatchdogFuture = Executors.newSingleThreadScheduledExecutor().scheduleAtFixedRate(
                this::checkAssignments, 10, 10, TimeUnit.SECONDS);
    }

    private void stopWatchdog() {
        partitionWatchdogFuture.cancel(true);
    }

    class ConsumerPartitionLevels {
        int lastKnownCount = -1;
        Instant lastChangeDate = Instant.ofEpochSecond(0);

        @Override
        public boolean equals(Object o) {
            if (this == o) return true;
            if (o == null || getClass() != o.getClass()) return false;
            ConsumerPartitionLevels that = (ConsumerPartitionLevels) o;
            return lastKnownCount == that.lastKnownCount &&
                    Objects.equals(lastChangeDate, that.lastChangeDate);
        }

        @Override
        public int hashCode() {
            return Objects.hash(lastKnownCount, lastChangeDate);
        }
    }

}

@@ -379,149 +379,9 @@
 ent.
-CopyOnWriteArraySet;%0Aimport java.util.concurrent.Executors;%0Aimport java.util.concurrent.ScheduledFuture;%0Aimport java.util.concurrent.TimeUnit
+*
 ;%0A%0A@
@@ -824,16 +824,64 @@
 gFuture;
+%0A    private ScheduledExecutorService scheduler;
 %0A%0A    pu
@@ -3067,39 +3067,25 @@
         
-partitionWatchdogFuture
+scheduler
  = Execu
@@ -3123,16 +3123,61 @@
 ecutor()
+;%0A        partitionWatchdogFuture = scheduler
 .schedul
@@ -3340,24 +3340,129 @@
 ncel(true);%0A
+        scheduler.shutdown();%0A        logger.debug(%22Partition assignment watchdog scheduler shut down%22);%0A
     %7D%0A%0A    c

Fix partition watchdog executor thread leak