Repository: apache/kafka
File: metadata/src/test/java/org/apache/kafka/controller/ClusterControlManagerTest.java (Java)

Code before:
            setSessionTimeoutNs(TimeUnit.MILLISECONDS.toNanos(brokerSessionTimeoutMs)).
            setTime(time).
            build();
        // first broker is fenced
        clusterControl.replay(new RegisterBrokerRecord().
            setBrokerEpoch(100).
            setBrokerId(0).
            setLogDirs(List.of(Uuid.fromString("Mj3CW3OSRi29cFeNJlXuAQ"))).
            setFenced(true).
            setInControlledShutdown(false), 10002);
        // second broker is in controlled shutdown
        clusterControl.replay(new RegisterBrokerRecord().
            setBrokerEpoch(200).
            setBrokerId(1).
            setLogDirs(List.of(Uuid.fromString("TyNK6XSSQJaJc2q9uflNHg"))).
            setFenced(false).
            setInControlledShutdown(true), 20002);
        clusterControl.activate();
        clusterControl.heartbeatManager().maybeUpdateControlledShutdownOffset(1, 20002);

        assertEquals(OptionalLong.of(1000L), clusterControl.heartbeatManager().tracker().
            contactTime(new BrokerIdAndEpoch(0, 100)));
        assertEquals(OptionalLong.of(1000L), clusterControl.heartbeatManager().tracker().
            contactTime(new BrokerIdAndEpoch(1, 200)));

        // even if session is still valid for old brokers, new registrations will succeed
        time.sleep(brokerSessionTimeoutMs / 2);
        clusterControl.registerBroker(new BrokerRegistrationRequestData().
                setClusterId("pjvUwj3ZTEeSVQmUiH3IJw").
                setBrokerId(0).
                setLogDirs(List.of(Uuid.fromString("yJGxmjfbQZSVFAlNM3uXZg"))).
                setFeatures(new BrokerRegistrationRequestData.FeatureCollection(
                    Set.of(new BrokerRegistrationRequestData.Feature().
                        setName(MetadataVersion.FEATURE_NAME).
                        setMinSupportedVersion(MetadataVersion.MINIMUM_VERSION.featureLevel()).
                        setMaxSupportedVersion(MetadataVersion.LATEST_PRODUCTION.featureLevel())).iterator())).
                setIncarnationId(Uuid.fromString("0H4fUu1xQEKXFYwB1aBjhg")),
            101L,
            finalizedFeatures,
            false);
        clusterControl.registerBroker(new BrokerRegistrationRequestData().
                setClusterId("pjvUwj3ZTEeSVQmUiH3IJw").
                setBrokerId(1).
                setLogDirs(List.of(Uuid.fromString("b66ybsWIQoygs01vdjH07A"))).
                setFeatures(new BrokerRegistrationRequestData.FeatureCollection(
                    Set.of(new BrokerRegistrationRequestData.Feature().
                        setName(MetadataVersion.FEATURE_NAME).
                        setMinSupportedVersion(MetadataVersion.MINIMUM_VERSION.featureLevel()).
                        setMaxSupportedVersion(MetadataVersion.LATEST_PRODUCTION.featureLevel())).iterator())).
                setIncarnationId(Uuid.fromString("vZKYST0pSA2HO5x_6hoO2Q")),
            201L,

Review comment (suggestion):
nit: I think the wording here should be something along the lines of:
> new registrations should succeed if a broker is fenced or in controlled shutdown, even if the last heartbeat was within the session timeout

Code after:
            setSessionTimeoutNs(TimeUnit.MILLISECONDS.toNanos(brokerSessionTimeoutMs)).
            setTime(time).
            build();
        // first broker is fenced
        clusterControl.replay(new RegisterBrokerRecord().
            setBrokerEpoch(100).
            setBrokerId(0).
            setLogDirs(List.of(Uuid.fromString("Mj3CW3OSRi29cFeNJlXuAQ"))).
            setFenced(true).
            setInControlledShutdown(false), 10002);
        // second broker is in controlled shutdown
        clusterControl.replay(new RegisterBrokerRecord().
            setBrokerEpoch(200).
            setBrokerId(1).
            setLogDirs(List.of(Uuid.fromString("TyNK6XSSQJaJc2q9uflNHg"))).
            setFenced(false).
            setInControlledShutdown(true), 20002);
        clusterControl.activate();
        clusterControl.heartbeatManager().maybeUpdateControlledShutdownOffset(1, 20002);

        assertEquals(OptionalLong.of(1000L), clusterControl.heartbeatManager().tracker().
            contactTime(new BrokerIdAndEpoch(0, 100)));
        assertEquals(OptionalLong.of(1000L), clusterControl.heartbeatManager().tracker().
            contactTime(new BrokerIdAndEpoch(1, 200)));

        time.sleep(brokerSessionTimeoutMs / 2);
        assertThrows(DuplicateBrokerRegistrationException.class, () ->
            clusterControl.registerBroker(new BrokerRegistrationRequestData().
                    setClusterId("pjvUwj3ZTEeSVQmUiH3IJw").
                    setBrokerId(0).
                    setLogDirs(List.of(Uuid.fromString("yJGxmjfbQZSVFAlNM3uXZg"))).
                    setFeatures(new BrokerRegistrationRequestData.FeatureCollection(
                        Set.of(new BrokerRegistrationRequestData.Feature().
                            setName(MetadataVersion.FEATURE_NAME).
                            setMinSupportedVersion(MetadataVersion.MINIMUM_VERSION.featureLevel()).
                            setMaxSupportedVersion(MetadataVersion.LATEST_PRODUCTION.featureLevel())).iterator())).
                    setIncarnationId(Uuid.fromString("0H4fUu1xQEKXFYwB1aBjhg")),
                101L,
                finalizedFeatures,
                false));
        // new registration should succeed immediatelly only if the broker is in controlled shutdown, 
        // even if the last heartbeat was within the session timeout 
        clusterControl.registerBroker(new BrokerRegistrationRequestData().
                setClusterId("pjvUwj3ZTEeSVQmUiH3IJw").
                setBrokerId(1).
                setLogDirs(List.of(Uuid.fromString("b66ybsWIQoygs01vdjH07A"))).
                setFeatures(new BrokerRegistrationRequestData.FeatureCollection(
                    Set.of(new BrokerRegistrationRequestData.Feature().
                        setName(MetadataVersion.FEATURE_NAME).
                        setMinSupportedVersion(MetadataVersion.MINIMUM_VERSION.featureLevel()).
                        setMaxSupportedVersion(MetadataVersion.LATEST_PRODUCTION.featureLevel())).iterator())).