@Component
public class PartitionObserver implements Consumer<Message<String>> {

  private static final Logger log = LoggerFactory.getLogger(PartitionObserver.class);

  private final Set<Integer> observedPartitions = new HashSet<>();
  private final AtomicLong messageCounter = new AtomicLong();

  @Override
  public void accept(Message<String> message) {
    var partition = message.getHeaders().get(KafkaHeaders.RECEIVED_PARTITION, Integer.class);
    log.info("Received message {} on partition {}", message.getPayload(), partition);
    observedPartitions.add(partition);
    messageCounter.incrementAndGet();
  }

  public long messageCount() {
    return messageCounter.get();
  }

  public int uniquePartitions() {
    return observedPartitions.size();
  }

  public void reset() {
    observedPartitions.clear();
    messageCounter.set(0);
  }
}

@SpringBootApplication
public class KafkaPartitioningIssueApplication implements CommandLineRunner {

  private static final Logger log =
      LoggerFactory.getLogger(KafkaPartitioningIssueApplication.class);

  @Autowired private PartitionObserver partitionObserver;
  @Autowired private StreamBridge streamBridge;
  @Autowired private ApplicationEventPublisher eventPublisher;

  @Value("{spring.application.name}")
  private String applicationName;

  public static void main(String[] args) {
    SpringApplication.run(KafkaPartitioningIssueApplication.class, args);
  }

  @Override
  public void run(String... args) throws Exception {
    int messageCount = 8;

    // initial batch
    processBatch(messageCount);

    // refresh
    partitionObserver.reset();
    eventPublisher.publishEvent(
        new RefreshRemoteApplicationEvent(this, applicationName, applicationName + ":**"));

    // second batch
    processBatch(messageCount);
  }

  private void processBatch(int messageCount) throws InterruptedException {
    for (int i = 0; i < messageCount; ++i) {
      streamBridge.send(
          "outbound",
          MessageBuilder.withPayload(String.valueOf(i)).setHeader("partitioning-key", i).build());
    }
    // wait for all messages to be consumed back
    while (partitionObserver.messageCount() < messageCount) {
      Thread.sleep(Duration.ofMillis(10));
    }
    log.info("Observed partitions: {}", partitionObserver.uniquePartitions());
  }
}

@Override
protected MessageHandler createProducerMessageHandler(
	final ProducerDestination destination,
	ExtendedProducerProperties<KafkaProducerProperties> producerProperties,
	MessageChannel channel, MessageChannel errorChannel) throws Exception {
// ...
Collection<PartitionInfo> partitions = provisioningProvider.getPartitionInfoForProducer(
		destination.getName(), producerFB, producerProperties);
this.topicsInUse.put(destination.getName(),
		new TopicInformation(null, partitions, false));
if (producerProperties.isPartitioned()
		&& producerProperties.getPartitionCount() < partitions.size()) {
	if (this.logger.isInfoEnabled()) {
		this.logger.info("The `partitionCount` of the producer for topic "
				+ destination.getName() + " is "
				+ producerProperties.getPartitionCount()
				+ ", smaller than the actual partition count of "
				+ partitions.size()
				+ " for the topic. The larger number will be used instead.");
	}
	producerProperties.setPartitionCount(partitions.size()); // <-- here
	List<ChannelInterceptor> interceptors = ((InterceptableChannel) channel)
			.getInterceptors();
	interceptors.forEach((interceptor) -> {
		if (interceptor instanceof PartitioningInterceptor partitioningInterceptor) {
			partitioningInterceptor.setPartitionCount(partitions.size());
		}
		else if (interceptor instanceof DefaultPartitioningInterceptor defaultPartitioningInterceptor) {
			defaultPartitioningInterceptor.setPartitionCount(partitions.size());
		}
	});
}
// ...