package org.apache.kafka.clients.producer.internals;
.....

public class DefaultPartitioner implements Partitioner {
    // Map ConcurrentMap để lưu trữ bộ đếm count cho từng topic 
    private final ConcurrentMap<String, AtomicInteger> topicCounterMap = new ConcurrentHashMap();

    // Constructor mặc định
    public DefaultPartitioner() {
    }

    // Phương thức cấu hình, hiện tại không làm gì
    public void configure(Map<String, ?> configs) {
    }

    // Phương thức xác định phân vùng cho một bản ghi
    public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
        // Lấy danh sách các phân vùng cho chủ đề
        List<PartitionInfo> partitions = cluster.partitionsForTopic(topic);
        int numPartitions = partitions.size();
        
        // Nếu không có khóa, sử dụng bộ đếm tuần tự để chọn phân vùng
        if (keyBytes == null) {
            int nextValue = this.nextValue(topic);
            // lấy về danh sách partition khả dụng để ghi. (Không bao gồm các partition off)
            List<PartitionInfo> availablePartitions = cluster.availablePartitionsForTopic(topic);
            
            // Nếu có phân vùng khả dụng, chọn một trong số chúng
            if (availablePartitions.size() > 0) {
                // Chuyển giá trị bộ đếm thành số dương và lấy phần dư khi chia cho số phân vùng khả dụng
                // Ví dụ topic có 5 partition.
                // part = 2107121300 % 5 ==> 0
                // part = 2107121301 % 5 ==> 1
                // part = 2107121302 % 5 ==> 2
                // part = 2107121303 % 5 ==> 3
                // part = 2107121304 % 5 ==> 4
                int part = Utils.toPositive(nextValue) % availablePartitions.size();
                return ((PartitionInfo)availablePartitions.get(part)).partition();
            } else {
                // Nếu không có phân vùng khả dụng, chọn ngẫu nhiên trong tất cả các phân vùng
                return Utils.toPositive(nextValue) % numPartitions;
            }
        } else {
            // Nếu có khóa, sử dụng hàm băm Murmur2 để chọn phân vùng
            return Utils.toPositive(Utils.murmur2(keyBytes)) % numPartitions;
        }
    }

    // Phương thức lấy giá trị tiếp theo của bộ đếm cho một chủ đề
    private int nextValue(String topic) {
        // Lấy bộ đếm cho chủ đề từ bản đồ
        AtomicInteger counter = (AtomicInteger)this.topicCounterMap.get(topic);
        
        // Nếu bộ đếm chưa tồn tại, khởi tạo nó với giá trị ngẫu nhiên
        if (null == counter) {
            counter = new AtomicInteger(ThreadLocalRandom.current().nextInt());
            AtomicInteger currentCounter = (AtomicInteger)this.topicCounterMap.putIfAbsent(topic, counter);
            if (currentCounter != null) {
                counter = currentCounter;
            }
        }

        // Tăng giá trị bộ đếm và trả về giá trị trước khi tăng
        return counter.getAndIncrement();
    }

    // Phương thức đóng, hiện tại không làm gì
    public void close() {
    }
}

package org.apache.kafka.clients.producer.internals;
...
public class DefaultPartitioner implements Partitioner {
    private final StickyPartitionCache stickyPartitionCache = new StickyPartitionCache();

    public DefaultPartitioner() {
    }

    public void configure(Map<String, ?> configs) {
    }

    public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster) {
        return this.partition(topic, key, keyBytes, value, valueBytes, cluster, cluster.partitionsForTopic(topic).size());
    }

    public int partition(String topic, Object key, byte[] keyBytes, Object value, byte[] valueBytes, Cluster cluster, int numPartitions) {
        return keyBytes == null 
        // nếu key null sử dụng stickyPartitionCache để lấy về partition 
        ? this.stickyPartitionCache.partition(topic, cluster) 
        // nếu key không null, tính hash của key và chia lấy dư cho numPartitions, 
        // trong case này nếu cùng một key thì sẽ luôn trả về cùng một partition
        : Utils.toPositive(Utils.murmur2(keyBytes)) % numPartitions;
    }

    public void close() {
    }

    public void onNewBatch(String topic, Cluster cluster, int prevPartition) {
        // Mỗi khi tạo mới Batch cho topic sẽ thực hiện tính toán lại partition cho batch mới.
        this.stickyPartitionCache.nextPartition(topic, cluster, prevPartition);
    }
}



package org.apache.kafka.clients.producer.internals;
...
// Lớp chịu trách nhiệm quản lý sticky partitioning cho các topic trong Kafka.
public class StickyPartitionCache {
    // Cache để lưu trữ chỉ số partition hiện tại cho mỗi topic.
    private final ConcurrentMap<String, Integer> indexCache = new ConcurrentHashMap<>();

    // Constructor: Khởi tạo một đối tượng StickyPartitionCache mới.
    public StickyPartitionCache() {
    }

    // Phương thức để lấy partition hiện tại cho một topic cụ thể.
    // Nếu partition đã được cache, trả về nó; nếu không, tính toán partition tiếp theo.
    // Đây là method được gọi bởi method partition trong class DefaultPartitioner để lấy về partititon
    public int partition(String topic, Cluster cluster) {
        Integer part = this.indexCache.get(topic);
        return part == null ? this.nextPartition(topic, cluster, -1) : part;
    }

    // Phương thức xác định partition tiếp theo cho một topic cụ thể.
    // Nếu không có partition khả dụng hoặc partition giống như trước, chọn một partition ngẫu nhiên.
    public int nextPartition(String topic, Cluster cluster, int prevPartition) {
        // Lấy danh sách tất cả các partition cho topic.
        List<PartitionInfo> partitions = cluster.partitionsForTopic(topic);
        Integer oldPart = this.indexCache.get(topic);
        Integer newPart = oldPart;

        // Nếu partition cũ(Partition hiện tại trong cache) không phải là partition trước đó, trả về partition cũ từ cache.
        if (oldPart != null && oldPart != prevPartition) {
            return this.indexCache.get(topic);
        } else {
            // Lấy danh sách các partition khả dụng cho topic.(Không bao gồm các partition off)
            List<PartitionInfo> availablePartitions = cluster.availablePartitionsForTopic(topic);

            // Nếu không có partition khả dụng, chọn một partition ngẫu nhiên.
            if (availablePartitions.size() < 1) {
                Integer random = Utils.toPositive(ThreadLocalRandom.current().nextInt());
                newPart = random % partitions.size();
            } 
            // Nếu chỉ có một partition khả dụng, sử dụng partition đó.
            else if (availablePartitions.size() == 1) {
                newPart = availablePartitions.get(0).partition();
            } 
            // Nếu có nhiều partition khả dụng, chọn một partition ngẫu nhiên khác với partition cũ.
            else {
                while(newPart == null || newPart.equals(oldPart)) {
                    int random = Utils.toPositive(ThreadLocalRandom.current().nextInt());
                    newPart = availablePartitions.get(random % availablePartitions.size()).partition();
                }
            }

            // Cập nhật cache với partition mới.
            if (oldPart == null) {
                this.indexCache.putIfAbsent(topic, newPart);
            } else {
                this.indexCache.replace(topic, prevPartition, newPart);
            }

            // Trả về partition mới từ cache.
            return this.indexCache.get(topic);
        }
    }
}

package org.apache.kafka.clients.producer.internals;
...
public class Sender implements Runnable {
  private long sendProducerData(long now) {
    Cluster cluster = metadata.fetch();
    // lấy về danh sách node có batch đủ điều kiện để gửi (đầy batch hoặc đủ lingerMs)
    // get the list of partitions with data ready to send
    RecordAccumulator.ReadyCheckResult result = this.accumulator.ready(cluster, now);

    // if there are any partitions whose leaders are not known yet, force metadata update
    if (!result.unknownLeaderTopics.isEmpty()) {
      // The set of topics with unknown leader contains topics with leader election pending as well as
      // topics which may have expired. Add the topic again to metadata to ensure it is included
      // and request metadata update, since there are messages to send to the topic.
      for (String topic : result.unknownLeaderTopics)
        this.metadata.add(topic, now);

      log.debug("Requesting metadata update due to unknown leader topics from the batched records: {}",
        result.unknownLeaderTopics);
      this.metadata.requestUpdate();
    }
    // Xóa các Node không sẵn sàng để gửi
    // remove any nodes we aren't ready to send to
    Iterator<Node> iter = result.readyNodes.iterator();
    long notReadyTimeout = Long.MAX_VALUE;
    while(iter.hasNext()){
        Node node = iter.next();
        // kiểm tra các node là sẵn sàng nhận request
        if (!this.client.ready(node, now)) {
            // Nếu node không sẵn sàng, xóa nó khỏi danh sách. Các node slow có tỉ lệ sẵn sàng thấp nên hay bị xóa.
            iter.remove();
            notReadyTimeout = Math.min(notReadyTimeout, this.client.pollDelayMs(node, now));
    }
    }

    // create produce requests trên các node sẵn sàng.
   // drain sẽ thực hiện Rút các data trong batch ra
    Map<Integer, List<ProducerBatch>> batches = this.accumulator.drain(cluster, result.readyNodes, this.maxRequestSize, now);
  ...
    // thực hiện gửi Batch và xóa batch sau khi gửi thành công.
    sendProduceRequests(batches, now);
    
}

package org.apache.kafka.clients.producer.internals;

import ....

public class BuiltInPartitioner {
    private final Logger log;
    private final String topic;
    private final int stickybatch.size;
    private volatile PartitionLoadStats partitionLoadStats = null;
    private final AtomicReference<StickyPartitionInfo> stickyPartitionInfo 
      = new AtomicReference();
    public static volatile Supplier<Integer> mockRandom = null;

    public BuiltInPartitioner(LogContext logContext, String topic, int stickybatch.size) {
        this.log = logContext.logger(BuiltInPartitioner.class);
        this.topic = topic;
        if (stickybatch.size < 1) {
            throw new IllegalArgumentException("stickybatch.size must be >= 1 but got " + stickybatch.size);
        } else {
            // được config bởi "batch.size".
            // Giá trị mặc định của "batch.size" là 16384 
            // (https://github.com/apache/kafka/blob/4a485ddb71c844acc8bf241feda1fcbffc5ce9be/clients/src/main/java/org/apache/kafka/clients/producer/ProducerConfig.java#L384)
            this.stickybatch.size = stickybatch.size;
        }
    }

    private int nextPartition(Cluster cluster) {
        // mockRandom là test của dev, có thể bỏ qua
        // response của line code này là 1 số ngẫu nhiên
        int random = mockRandom != null ? (Integer)mockRandom.get() : Utils.toPositive(ThreadLocalRandom.current().nextInt());
        PartitionLoadStats partitionLoadStats = this.partitionLoadStats;
        int partition;
        //nếu partitionLoadStats == null sẽ lấy ngẫu nhiên giống 2.4.0 đến 3.2.3 
        if (partitionLoadStats == null) {
            List<PartitionInfo> availablePartitions = cluster.availablePartitionsForTopic(this.topic);
            if (!availablePartitions.isEmpty()) {
                // từ số random, chia lấy dư cho số lượng availablePartitions
                partition = ((PartitionInfo)availablePartitions.get(random % availablePartitions.size())).partition();
            } else {
                // nếu không có availablePartitions, lấy ngẫu nhiên ở tất cho partition
                List<PartitionInfo> partitions = cluster.partitionsForTopic(this.topic);
                partition = random % partitions.size();
            }
        } else {
            // tính toán dựa trên lượng tải của partition
            // Nếu lượng tải quá lớn cùng một thời gian ngắn, partitionLoadStats sẽ not null và sử dụng partitionLoadStats để tính toán partition mới
            // Sau 1 khoản thời gian lượng tải hết lớn, partitionLoadStats sẽ null và sử dụng random
            // 
            assert partitionLoadStats.length > 0;

            int[] cumulativeFrequencyTable = partitionLoadStats.cumulativeFrequencyTable;
            
            // tạo trọng số ngẫu nhiên bằng cách lấy dư của random / cho giá trị cuối cùng trong bảng tần suất tích lũy.
            int weightedRandom = random % cumulativeFrequencyTable[partitionLoadStats.length - 1];
            
            // Sử dụng tìm kiếm nhị phân để tìm chỉ số của partition tương ứng với số ngẫu nhiên có trọng số trong bảng tần suất tích lũy.
            // nếu không tìm thấy, result sẽ trả về vị trí giá trị nên được thêm.
            int searchResult = Arrays.binarySearch(cumulativeFrequencyTable, 0, partitionLoadStats.length, weightedRandom);
            int partitionIndex = Math.abs(searchResult + 1);

            assert partitionIndex < partitionLoadStats.length;

            partition = partitionLoadStats.partitionIds[partitionIndex];
        }

        this.log.trace("Switching to partition {} in topic {}", partition, this.topic);
        return partition;
    }

    void updatePartitionInfo(StickyPartitionInfo partitionInfo, int appendedBytes, Cluster cluster, boolean enableSwitch) {
        if (partitionInfo != null) {
            assert partitionInfo == this.stickyPartitionInfo.get();

            int producedBytes = partitionInfo.producedBytes.addAndGet(appendedBytes);
            if (producedBytes >= this.stickybatch.size * 2) {
                this.log.trace("Produced {} bytes, exceeding twice the batch size of {} bytes, with switching set to {}", new Object[]{producedBytes, this.stickybatch.size, enableSwitch});
            }

            if (producedBytes >= this.stickybatch.size && enableSwitch || producedBytes >= this.stickybatch.size * 2) {
                // Nếu tổng số producedBytes hiện tại lớn hớn config stickybatch.size và enableSwitch = true hoặc producedBytes lớn hơn gấp 2 lần stickybatch.size thì sẽ force update 
                StickyPartitionInfo newPartitionInfo = new StickyPartitionInfo(this.nextPartition(cluster));
                this.stickyPartitionInfo.set(newPartitionInfo);
            }

        }
    }

    private final static class PartitionLoadStats {
      // load tải table, có size array tối đa  = số lượng partition
      public final int[] cumulativeFrequencyTable;
      // danh sach partition id cua topic
      public final int[] partitionIds;
      public final int length;
      public PartitionLoadStats(int[] cumulativeFrequencyTable, int[] partitionIds, int length) {
        assert cumulativeFrequencyTable.length == partitionIds.length;
        assert length <= cumulativeFrequencyTable.length;
        this.cumulativeFrequencyTable = cumulativeFrequencyTable;
        this.partitionIds = partitionIds;
        this.length = length;
      }
    }
}

  static int estimatebatch.sizeUpperBound(ByteBuffer key, ByteBuffer value, Header[] headers) {
    // RECORD_BATCH_OVERHEAD = 61 ;
    return RECORD_BATCH_OVERHEAD + DefaultRecord.recordSizeUpperBound(key, value, headers);
  }

  static int recordSizeUpperBound(ByteBuffer key, ByteBuffer value, Header[] headers) {
    int keySize = key == null ? -1 : key.remaining();
    int valueSize = value == null ? -1 : value.remaining();
    // MAX_RECORD_OVERHEAD = 21
    return MAX_RECORD_OVERHEAD + sizeOf(keySize, valueSize, headers);
  }
  
  private static int sizeOf(int keySize, int valueSize, Header[] headers) {
    int size = 0;
    if (keySize < 0)
        size += NULL_VARINT_SIZE_BYTES;
    else
        size += ByteUtils.sizeOfVarint(keySize) + keySize;

    if (valueSize < 0)
        size += NULL_VARINT_SIZE_BYTES;
    else
        size += ByteUtils.sizeOfVarint(valueSize) + valueSize;

    if (headers == null)
        throw new IllegalArgumentException("Headers cannot be null");

    size += ByteUtils.sizeOfVarint(headers.length);
    for (Header header : headers) {
      String headerKey = header.key();
      if (headerKey == null)
        throw new IllegalArgumentException("Invalid null header key found in headers");
        int headerKeySize = Utils.utf8Length(headerKey);
        size += ByteUtils.sizeOfVarint(headerKeySize) + headerKeySize;
  
        byte[] headerValue = header.value();
        if (headerValue == null) {
            size += NULL_VARINT_SIZE_BYTES;
        } else {
            size += ByteUtils.sizeOfVarint(headerValue.length) + headerValue.length;
        }
      }
    return size;
   }

public class CalPartitionLoad {
  public static void main(String[] args) {
    Integer[] queueSizes = {8, 3, 14, 8, 5};
    run(queueSizes);
    // Output: [7, 19, 20, 27, 37]
  }

  public static void run(Integer[] queueSizes){
    // queueSizes length = max partition, queueSizes là số lượng batch data đang cho gui den partition
    /**
     * ví dụ : Integer[] queueSizes = {8, 3, 14};
     * Có 3 partiotn là 0,1,2
     * partiton 0 có 8 batch đang chờ gửi
     * partiton 1 có 3 batch đang chờ gửi
     * partiton 14 có 8 batch đang chờ gửi
     */
    // lấy về số lớn nhất + 1
    int maxSizePlus1 = queueSizes[0];

    // trong thực tế, length là tham số truyền vào biểu diễn số lượng partition đang có leader và có thể gửi data
    int length = queueSizes.length;

    boolean allEqual = true;
    for (int i = 1; i < length; i++) {
      if (queueSizes[i] != maxSizePlus1)
        allEqual = false;
      if (queueSizes[i] > maxSizePlus1)
        maxSizePlus1 = queueSizes[i];
    }
    ++maxSizePlus1;

    //
    if (allEqual) {
      System.out.println("allEqual && queueSizes.length == queueSizes.length");
      return;
    }
    // mỗi phần tử tiếp theo được tính bằng cách sử dụng giá trị lớn nhất và cộng dồn các phần tử trước đó.
    //Với queueSizes = {8, 3, 14} và maxSizePlus1 = 15:
    //queueSizes[0] = 15 - 8 = 7
    //queueSizes[1] = 15 - 3 + 7 = 19
    //queueSizes[2] = 15 - 14 + 19 = 20
    //Kết quả cuối cùng: queueSizes = {7, 19, 20}.
    queueSizes[0] = maxSizePlus1 - queueSizes[0];
    for (int i = 1; i < length; i++) {
      queueSizes[i] = maxSizePlus1 - queueSizes[i] + queueSizes[i - 1];
    }
    System.out.println("Output: "+ Arrays.stream(queueSizes).toList());
  }
}

@Slf4j
public class KeyNullLoopDelay {
  @SneakyThrows
  public static void main(String[] args) throws IOException {
    System.setProperty("org.slf4j.simpleLogger.defaultLogLevel", "ALL");
    final var props = new Properties();
    props.setProperty(ProducerConfig.CLIENT_ID_CONFIG, "java-producer-producerRecordPartition-KeyNotNull");
    props.setProperty(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:29091,localhost:29092,localhost:29093,localhost:29094");
    props.setProperty(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());
    props.setProperty(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());
    props.setProperty(ProducerConfig.BATCH_SIZE_CONFIG, "5000");
    // 15s
    props.setProperty(ProducerConfig.LINGER_MS_CONFIG, "15000");

    try (var producer = new KafkaProducer<Object, String>(props)) {
      while (true) {
        final var messageProducerRecord = new ProducerRecord<>(
          "topic-rep-1-partition-10",     //topic name
          // 36 byte
          UUID.randomUUID().toString()        // value
        );
        Integer numberSend = 113;
        for (int i = 1; i <= numberSend; i++) {
          producer.send(messageProducerRecord);
        }
        // 30s
        Thread.sleep(30000);
      }
    }
  }
}

@Slf4j
public class KeyNull {
  @SneakyThrows
  public static void main(String[] args) throws IOException {
    System.setProperty("org.slf4j.simpleLogger.defaultLogLevel", "ALL");
    final var props = new Properties();
    props.setProperty(ProducerConfig.CLIENT_ID_CONFIG, "java-producer-producerRecordPartition-KeyNotNull");
    props.setProperty(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, "localhost:29091,localhost:29092,localhost:29093,localhost:29094");
    props.setProperty(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());
    props.setProperty(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, StringSerializer.class.getName());
    props.setProperty(ProducerConfig.BATCH_SIZE_CONFIG, "5000");

    try (var producer = new KafkaProducer<Object, String>(props)) {
        final var messageProducerRecord = new ProducerRecord<>(
          "topic-rep-1-partition-10",     //topic name
          // 36 byte
          UUID.randomUUID().toString()        // value
        );
        Integer numberSend = 100_000;
        for (int i = 1; i <= numberSend; i++) {
          producer.send(messageProducerRecord);
        }
    }
  }
}