Repository: apache/kafka
File: group-coordinator/src/main/java/org/apache/kafka/coordinator/group/runtime/CoordinatorRuntime.java (Java)

Code before:
                    tp, partitionEpoch);
            }
        });
    }

    /**
     * A new metadata image is available.
     *
     * @param newImage  The new metadata image.
     * @param delta     The metadata delta.
     */
    public void onNewMetadataImage(
        MetadataImage newImage,
        MetadataDelta delta
    ) {
        throwIfNotRunning();
        log.debug("Scheduling applying of a new metadata image with offset {}.", newImage.offset());

        // Update global image.
        metadataImage = newImage;

        // Push an event for each coordinator.
        coordinators.keySet().forEach(tp -> {
            scheduleInternalOperation("UpdateImage(tp=" + tp + ", offset=" + newImage.offset() + ")", tp, () -> {
                CoordinatorContext context = coordinators.get(tp);
                if (context != null && context.state == CoordinatorState.ACTIVE) {
                    try {
                        context.lock.lock();
                        log.debug("Applying new metadata image with offset {} to {}.", newImage.offset(), tp);
                        context.coordinator.onNewMetadataImage(newImage, delta);
                    } finally {
                        context.lock.unlock();
                    }
                } else {
                    log.debug("Ignoring new metadata image with offset {} for {} because the coordinator is not active.",
                        newImage.offset(), tp);
                }
            });
        });
    }

    /**
     * Closes the runtime. This closes all the coordinators currently registered
     * in the runtime.
     *
     * @throws Exception
     */
    public void close() throws Exception {
        if (!isRunning.compareAndSet(true, false)) {
            log.warn("Coordinator runtime is already shutting down.");
            return;

Review comment (suggestion):
I noticed there are still some places we don't want the context to be null -- ie loading the partition. (And for all the run methods we want the coordinator to be active) Is the general rationale for this change that not having the conetext/active state is ok if we expect the coordinator moved? Is it possible we could try to load and the coordinator moves again? Just trying to understand the cases we accept null or non-active coordinators vs when we throw errors.

Code after:
        MetadataImage newImage,
        MetadataDelta delta
    ) {
        throwIfNotRunning();
        log.debug("Scheduling applying of a new metadata image with offset {}.", newImage.offset());

        // Update global image.
        metadataImage = newImage;

        // Push an event for each coordinator.
        coordinators.keySet().forEach(tp -> {
            scheduleInternalOperation("UpdateImage(tp=" + tp + ", offset=" + newImage.offset() + ")", tp, () -> {
                CoordinatorContext context = coordinators.get(tp);
                if (context != null) {
                    context.lock.lock();
                    try {
                        if (context.state == CoordinatorState.ACTIVE) {
                            // The new image can be applied to the coordinator only if the coordinator
                            // exists and is in the active state.
                            log.debug("Applying new metadata image with offset {} to {}.", newImage.offset(), tp);
                            context.coordinator.onNewMetadataImage(newImage, delta);
                        } else {
                            log.debug("Ignored new metadata image with offset {} for {} because the coordinator is not active.",
                                newImage.offset(), tp);
                        }
                    } finally {
                        context.lock.unlock();
                    }
                } else {
                    log.debug("Ignored new metadata image with offset {} for {} because the coordinator does not exist.",
                        newImage.offset(), tp);
                }
            });
        });
    }

    /**
     * Closes the runtime. This closes all the coordinators currently registered
     * in the runtime.
     *
     * @throws Exception
     */
    public void close() throws Exception {
        if (!isRunning.compareAndSet(true, false)) {
            log.warn("Coordinator runtime is already shutting down.");
            return;
        }

        log.info("Closing coordinator runtime.");
        Utils.closeQuietly(loader, "loader");
        Utils.closeQuietly(timer, "timer");