From f2c318342f4deb5c0087b3f79c47fbfcd4dac12d Mon Sep 17 00:00:00 2001 From: Sijie Guo Date: Mon, 22 Jul 2019 01:40:25 +0800 Subject: [PATCH 1/8] Kafka coordinator service (part 1): Add GroupMetadata and MemberMetadata *Motivation* Import Kafka coordinator implementation. This is the first part to introduce GroupMetadata and MemberMetadata. --- pom.xml | 2 + .../kop/coordinator/group/GroupConfig.java | 28 ++ .../coordinator/group/GroupCoordinator.java | 44 ++ .../kop/coordinator/group/GroupMetadata.java | 345 ++++++++++++++++ .../kop/coordinator/group/GroupState.java | 99 +++++ .../coordinator/group/JoinGroupResult.java | 33 ++ .../kop/coordinator/group/MemberMetadata.java | 169 ++++++++ .../kop/coordinator/group/OffsetConfig.java | 25 ++ .../kop/coordinator/group/package-info.java | 17 + .../coordinator/group/GroupMetadataTest.java | 390 ++++++++++++++++++ .../coordinator/group/MemberMetadataTest.java | 126 ++++++ 11 files changed, 1278 insertions(+) create mode 100644 src/main/java/io/streamnative/kop/coordinator/group/GroupConfig.java create mode 100644 src/main/java/io/streamnative/kop/coordinator/group/GroupCoordinator.java create mode 100644 src/main/java/io/streamnative/kop/coordinator/group/GroupMetadata.java create mode 100644 src/main/java/io/streamnative/kop/coordinator/group/GroupState.java create mode 100644 src/main/java/io/streamnative/kop/coordinator/group/JoinGroupResult.java create mode 100644 src/main/java/io/streamnative/kop/coordinator/group/MemberMetadata.java create mode 100644 src/main/java/io/streamnative/kop/coordinator/group/OffsetConfig.java create mode 100644 src/main/java/io/streamnative/kop/coordinator/group/package-info.java create mode 100644 src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataTest.java create mode 100644 src/test/java/io/streamnative/kop/coordinator/group/MemberMetadataTest.java diff --git a/pom.xml b/pom.xml index d60a845b56..18a5b49a88 100644 --- a/pom.xml +++ b/pom.xml @@ -249,7 +249,9 @@ ${javac.target} ${javac.target} + -Xlint:deprecation -Xlint:unchecked diff --git a/src/main/java/io/streamnative/kop/coordinator/group/GroupConfig.java b/src/main/java/io/streamnative/kop/coordinator/group/GroupConfig.java new file mode 100644 index 0000000000..d609de0cef --- /dev/null +++ b/src/main/java/io/streamnative/kop/coordinator/group/GroupConfig.java @@ -0,0 +1,28 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.coordinator.group; + +import lombok.Data; + +/** + * Group configuration. + */ +@Data +public class GroupConfig { + + private final int groupMinSessionTimeoutMs; + private final int groupMaxSessionTimeoutMs; + private final int groupInitialRebalanceDelayMs; + +} diff --git a/src/main/java/io/streamnative/kop/coordinator/group/GroupCoordinator.java b/src/main/java/io/streamnative/kop/coordinator/group/GroupCoordinator.java new file mode 100644 index 0000000000..ab631869b6 --- /dev/null +++ b/src/main/java/io/streamnative/kop/coordinator/group/GroupCoordinator.java @@ -0,0 +1,44 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.coordinator.group; + +import io.streamnative.kop.coordinator.group.GroupMetadata.GroupSummary; +import java.util.Collections; + +/** + * Group coordinator. + */ +public class GroupCoordinator { + + static final String NoState = ""; + static final String NoProtocolType = ""; + static final String NoProtocol = ""; + static final String NoLeader = ""; + static final int NoGeneration = -1; + static final String NoMemberId = ""; + static final GroupSummary EmptyGroup = new GroupSummary( + NoState, + NoProtocolType, + NoProtocol, + Collections.emptyList() + ); + static final GroupSummary DeadGroup = new GroupSummary( + GroupState.Dead.toString(), + NoProtocolType, + NoProtocol, + Collections.emptyList() + ); + + +} diff --git a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadata.java b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadata.java new file mode 100644 index 0000000000..aca6780093 --- /dev/null +++ b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadata.java @@ -0,0 +1,345 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.coordinator.group; + +import static com.google.common.base.Preconditions.checkArgument; +import static com.google.common.base.Preconditions.checkState; + +import com.google.common.base.MoreObjects; +import com.google.common.base.MoreObjects.ToStringHelper; +import com.google.common.collect.Sets; +import io.streamnative.kop.coordinator.group.MemberMetadata.MemberSummary; +import java.util.Collections; +import java.util.Comparator; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Map.Entry; +import java.util.Objects; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.locks.ReentrantLock; +import java.util.stream.Collectors; +import javax.annotation.concurrent.NotThreadSafe; +import lombok.Data; +import lombok.Setter; +import lombok.experimental.Accessors; +import org.apache.commons.lang3.StringUtils; + +/** + * Group contains the following metadata: + * + *

Membership metadata: + * 1. Members registered in this group + * 2. Current protocol assigned to the group (e.g. partition assignment strategy for consumers) + * 3. Protocol metadata associated with group members + * + *

State metadata: + * 1. group state + * 2. generation id + * 3. leader id + */ +@NotThreadSafe +@Setter +@Accessors(fluent = true) +class GroupMetadata { + + private static final Map> validPreviousStates = new HashMap<>(); + + static { + validPreviousStates.put( + GroupState.Dead, + Sets.newHashSet( + GroupState.Stable, + GroupState.PreparingRebalance, + GroupState.CompletingRebalance, + GroupState.Empty, + GroupState.Dead + ) + ); + + validPreviousStates.put( + GroupState.CompletingRebalance, + Sets.newHashSet( + GroupState.PreparingRebalance + ) + ); + + validPreviousStates.put( + GroupState.Stable, + Sets.newHashSet( + GroupState.CompletingRebalance + ) + ); + + validPreviousStates.put( + GroupState.PreparingRebalance, + Sets.newHashSet( + GroupState.Stable, + GroupState.CompletingRebalance, + GroupState.Empty + ) + ); + + validPreviousStates.put( + GroupState.Empty, + Sets.newHashSet( + GroupState.PreparingRebalance + ) + ); + } + + public static GroupMetadata loadGroup( + String groupId, + GroupState initialState, + int generationId, + String protocolType, + String protocol, + String leaderId, + Iterable members + ) { + GroupMetadata metadata = new GroupMetadata(groupId, initialState) + .generationId(generationId) + .protocolType( + StringUtils.isEmpty(protocolType) ? Optional.empty() : Optional.of(protocolType) + ) + .protocol(Optional.ofNullable(protocol)) + .leaderId(Optional.ofNullable(leaderId)); + members.forEach(metadata::add); + return metadata; + } + + /** + * Class used to represent group metadata for the ListGroups API. + */ + @Data + static class GroupOverview { + private final String groupId; + private final String protocolType; + } + + /** + * Class used to represent group metadata for the DescribeGroup API. + */ + @Data + static class GroupSummary { + private final String state; + private final String protocolType; + private final String protocol; + private final List members; + } + + private final String groupId; + private final ReentrantLock lock = new ReentrantLock(); + private GroupState state; + + private Optional protocolType = Optional.empty(); + private long generationId = 0L; + private Optional leaderId = Optional.empty(); + private Optional protocol = Optional.empty(); + + // state management + private final Map members = new HashMap<>(); + + GroupMetadata(String groupId, GroupState initialState) { + this.groupId = groupId; + this.state = initialState; + } + + public GroupState currentState() { + return state; + } + + public long generationId() { + return generationId; + } + + public boolean is(GroupState groupState) { + return state == groupState; + } + + public boolean not(GroupState groupState) { + return state != groupState; + } + + public boolean has(String memberId) { + return members.containsKey(memberId); + } + + public boolean isLeader(String memberId) { + return Objects.equals(leaderId.orElse(null), memberId); + } + + public String leaderOrNull() { + return leaderId.orElse(null); + } + + public String protocolOrNull() { + return protocol.orElse(null); + } + + public List notYetRejoinedMembers() { + return members.values() + .stream() + .filter(e -> e.awaitingJoinCallback() == null) + .collect(Collectors.toList()); + } + + private Set candidateProtocols() { + return members.values().stream() + .map(MemberMetadata::protocols) + .reduce((p1, p2) -> { + Set newProtocols = new HashSet<>(); + newProtocols.addAll(Sets.intersection(p1, p2)); + return newProtocols; + }) + .orElse(Collections.emptySet()); + } + + public boolean supportsProtocols(Set memberProtocols) { + return members.isEmpty() + || !Sets.intersection(memberProtocols, candidateProtocols()).isEmpty(); + } + + public void initNextGeneration() { + checkArgument(notYetRejoinedMembers().isEmpty()); + if (!members.isEmpty()) { + generationId += 1; + protocol = Optional.ofNullable(selectProtocol()); + transitionTo(GroupState.CompletingRebalance); + } else { + generationId += 1; + protocol = Optional.empty(); + transitionTo(GroupState.Empty); + } + } + + public void add(MemberMetadata member) { + if (members.isEmpty()) { + this.protocolType = Optional.of(member.protocolType()); + } + + checkArgument(groupId == member.groupId()); + checkArgument(Objects.equals(protocolType.orElse(null), member.protocolType())); + checkArgument(supportsProtocols(member.protocols())); + + if (!leaderId.isPresent()) { + leaderId = Optional.of(member.memberId()); + } + + members.put(member.memberId(), member); + } + + public void remove(String memberId) { + members.remove(memberId); + if (isLeader(memberId)) { + if (members.isEmpty()) { + leaderId = Optional.empty(); + } else { + leaderId = members.keySet().stream().findFirst(); + } + } + } + + public boolean canReblance() { + return validPreviousStates.get(GroupState.PreparingRebalance).contains(state); + } + + public void transitionTo(GroupState groupState) { + assertValidTransition(groupState); + state = groupState; + } + + private void assertValidTransition(GroupState targetState) { + if (!validPreviousStates.get(targetState).contains(state)) { + throw new IllegalStateException(("Group %s should be in the %s states before moving" + + " to %s state. Instead it is in %s state" + ).format( + groupId, + StringUtils.join(validPreviousStates.get(targetState), ","), + targetState, + state)); + } + } + + public String selectProtocol() { + checkState( + !members.isEmpty(), + "Cannot select protocol for empty group"); + + Set candidates = candidateProtocols(); + + return members.values().stream() + .map(m -> m.vote(candidates)) + .collect(Collectors.groupingBy(protocol -> protocol)) + .entrySet() + .stream() + .max(Comparator.comparingInt(o -> o.getValue().size())) + .map(Entry::getKey) + .orElse(null); + } + + public GroupSummary summary() { + if (is(GroupState.Stable)) { + String protocol = protocolOrNull(); + checkState( + protocol != null, + "Invalid null group protocol for stable group"); + + List summaries = members.values() + .stream() + .map(member -> member.summary(protocol)) + .collect(Collectors.toList()); + + return new GroupSummary( + state.toString(), + protocolType.orElse(""), + protocol, + summaries + ); + } else { + List summaries = members.values() + .stream() + .map(member -> member.summaryNoMetadata()) + .collect(Collectors.toList()); + + return new GroupSummary( + state.toString(), + protocolType.orElse(""), + GroupCoordinator.NoProtocol, + summaries + ); + } + } + + public GroupOverview overview() { + return new GroupOverview( + groupId, + protocolType.orElse("") + ); + } + + @Override + public String toString() { + ToStringHelper helper = MoreObjects.toStringHelper("GroupMetadata") + .add("groupId", groupId) + .add("generation", generationId) + .add("protocolType", protocolType) + .add("state", state) + .add("members", members); + return helper.toString(); + } + +} diff --git a/src/main/java/io/streamnative/kop/coordinator/group/GroupState.java b/src/main/java/io/streamnative/kop/coordinator/group/GroupState.java new file mode 100644 index 0000000000..037f9c61b2 --- /dev/null +++ b/src/main/java/io/streamnative/kop/coordinator/group/GroupState.java @@ -0,0 +1,99 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.coordinator.group; + +/** + * The state of the group. + * + *

This class is rewritten following Kafka's logic. + */ +public enum GroupState { + + /** + * Group is preparing to rebalance. + * + *

action: respond to heartbeats with REBALANCE_IN_PROGRESS + * respond to sync group with REBALANCE_IN_PROGRESS + * remove member on leave group request + * park join group requests from new or existing members until all expected members have joined + * allow offset commits from previous generation + * allow offset fetch requests + * transition: some members have joined by the timeout => CompletingRebalance + * all members have left the group => Empty + * group is removed by partition emigration => Dead + */ + PreparingRebalance, + + /** + * Group is awaiting state assignment from the leader. + * + *

action: respond to heartbeats with REBALANCE_IN_PROGRESS + * respond to offset commits with REBALANCE_IN_PROGRESS + * park sync group requests from followers until transition to Stable + * allow offset fetch requests + * transition: sync group with state assignment received from leader => Stable + * join group from new member or existing member with updated metadata => PreparingRebalance + * leave group from existing member => PreparingRebalance + * member failure detected => PreparingRebalance + * group is removed by partition emigration => Dead + */ + CompletingRebalance, + + /** + * Group is stable. + * + *

action: respond to member heartbeats normally + * respond to sync group from any member with current assignment + * respond to join group from followers with matching metadata with current group metadata + * allow offset commits from member of current generation + * allow offset fetch requests + * transition: member failure detected via heartbeat => PreparingRebalance + * leave group from existing member => PreparingRebalance + * leader join-group received => PreparingRebalance + * follower join-group with new metadata => PreparingRebalance + * group is removed by partition emigration => Dead + */ + Stable, + + /** + * Group has no more members and its metadata is being removed. + * + *

action: respond to join group with UNKNOWN_MEMBER_ID + * respond to sync group with UNKNOWN_MEMBER_ID + * respond to heartbeat with UNKNOWN_MEMBER_ID + * respond to leave group with UNKNOWN_MEMBER_ID + * respond to offset commit with UNKNOWN_MEMBER_ID + * allow offset fetch requests + * transition: Dead is a final state before group metadata is cleaned up, so there are no transitions + */ + Dead, + + /** + * Group has no more members, but lingers until all offsets have expired. This state + * also represents groups which use Kafka only for offset commits and have no members. + * + *

action: respond normally to join group from new members + * respond to sync group with UNKNOWN_MEMBER_ID + * respond to heartbeat with UNKNOWN_MEMBER_ID + * respond to leave group with UNKNOWN_MEMBER_ID + * respond to offset commit with UNKNOWN_MEMBER_ID + * allow offset fetch requests + * transition: last offsets removed in periodic expiration task => Dead + * join group from a new member => PreparingRebalance + * group is removed by partition emigration => Dead + * group is removed by expiration => Dead + */ + Empty + +} diff --git a/src/main/java/io/streamnative/kop/coordinator/group/JoinGroupResult.java b/src/main/java/io/streamnative/kop/coordinator/group/JoinGroupResult.java new file mode 100644 index 0000000000..ae7bbfa4da --- /dev/null +++ b/src/main/java/io/streamnative/kop/coordinator/group/JoinGroupResult.java @@ -0,0 +1,33 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.coordinator.group; + +import java.util.Map; +import lombok.Data; +import org.apache.kafka.common.protocol.Errors; + +/** + * The result of a join group operation. + */ +@Data +public class JoinGroupResult { + + private final Map members; + private final String memberId; + private final int generationId; + private final String subProtocol; + private final String leaderId; + private final Errors error; + +} diff --git a/src/main/java/io/streamnative/kop/coordinator/group/MemberMetadata.java b/src/main/java/io/streamnative/kop/coordinator/group/MemberMetadata.java new file mode 100644 index 0000000000..bb6f728768 --- /dev/null +++ b/src/main/java/io/streamnative/kop/coordinator/group/MemberMetadata.java @@ -0,0 +1,169 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.coordinator.group; + +import static com.google.common.base.Preconditions.checkArgument; + +import com.google.common.base.MoreObjects; +import com.google.common.base.MoreObjects.ToStringHelper; +import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; +import java.util.Arrays; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.function.BiConsumer; +import java.util.function.Consumer; +import javax.annotation.concurrent.NotThreadSafe; +import lombok.Data; +import lombok.Getter; +import lombok.RequiredArgsConstructor; +import lombok.Setter; +import lombok.experimental.Accessors; +import org.apache.kafka.common.protocol.Errors; + +/** + * Member metadata contains the following metadata: + * + *

Heartbeat metadata: + * 1. negotiated heartbeat session timeout + * 2. timestamp of the latest heartbeat + * + *

Protocol metadata: + * 1. the list of supported protocols (ordered by preference) + * 2. the metadata associated with each protocol + * + *

In addition, it also contains the following state information: + * + *

1. Awaiting rebalance callback: when the group is in the prepare-rebalance state, + * its rebalance callback will be kept in the metadata if the + * member has sent the join group request + * 2. Awaiting sync callback: when the group is in the awaiting-sync state, its sync callback + * is kept in metadata until the leader provides the group assignment + * and the group transitions to stable + */ +@RequiredArgsConstructor +@NotThreadSafe +@Accessors(fluent = true) +@Getter +@Setter +@SuppressFBWarnings({ + "EI_EXPOSE_REP", + "EI_EXPOSE_REP2" +}) +public class MemberMetadata { + + /** + * Summary of member metadata. + */ + @Data + @SuppressFBWarnings({ + "EI_EXPOSE_REP", + "EI_EXPOSE_REP2" + }) + public static class MemberSummary { + private final String memberId; + private final String clientId; + private final String clientHost; + private final byte[] metadata; + private final byte[] assignment; + } + + private final String memberId; + private final String groupId; + private final String clientId; + private final String clientHost; + private final int rebalanceTimeoutMs; + private final int sessionTimeoutMs; + private final String protocolType; + private final Map supportedProtocols; + + private byte[] assignment = new byte[0]; + private Consumer awaitingJoinCallback = null; + private BiConsumer awaitingSyncCallback = null; + private long latestHeartbeat = -1L; + private boolean isLeaving = false; + + public Set protocols() { + return supportedProtocols.keySet(); + + } + + public byte[] metadata(String protocol) { + byte[] metadata = supportedProtocols.get(protocol); + checkArgument(metadata != null, "Member does not support protocol"); + return metadata; + } + + /** + * Check if the provided protocol metadata matches the currently stored metadata. + */ + public boolean matches(Map protocols) { + if (protocols.size() != this.supportedProtocols.size()) { + return false; + } + + for (Map.Entry protocolEntry : protocols.entrySet()) { + byte[] p1 = protocolEntry.getValue(); + byte[] p2 = this.supportedProtocols.get(protocolEntry.getKey()); + if (p2 == null || !Arrays.equals(p1, p2)) { + return false; + } + } + return true; + } + + public MemberSummary summary(String protocol) { + return new MemberSummary( + memberId, + clientId, + clientHost, + metadata(protocol), + assignment + ); + } + + public MemberSummary summaryNoMetadata() { + return new MemberSummary( + memberId, + clientId, + clientHost, + new byte[0], + new byte[0] + ); + } + + public String vote(Set candidates) { + Optional> voteProtocol = this.supportedProtocols.entrySet() + .stream() + .filter(p -> candidates.contains(p.getKey())) + .findFirst(); + checkArgument( + voteProtocol.isPresent(), + "Member does not support any of the candidate protocols"); + return voteProtocol.get().getKey(); + } + + @Override + public String toString() { + ToStringHelper helper = MoreObjects.toStringHelper("MemberMetadata") + .add("memberId", memberId) + .add("clientId", clientId) + .add("clientHost", clientHost) + .add("sessionTimeoutMs", sessionTimeoutMs) + .add("rebalanceTimeoutMs", rebalanceTimeoutMs) + .add("supportedProtocols", protocols().stream()); + return helper.toString(); + } + +} diff --git a/src/main/java/io/streamnative/kop/coordinator/group/OffsetConfig.java b/src/main/java/io/streamnative/kop/coordinator/group/OffsetConfig.java new file mode 100644 index 0000000000..c90f02b64c --- /dev/null +++ b/src/main/java/io/streamnative/kop/coordinator/group/OffsetConfig.java @@ -0,0 +1,25 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.coordinator.group; + +import lombok.Builder; +import lombok.Data; + +/** + * Offset configuration. + */ +@Builder +@Data +public class OffsetConfig { +} diff --git a/src/main/java/io/streamnative/kop/coordinator/group/package-info.java b/src/main/java/io/streamnative/kop/coordinator/group/package-info.java new file mode 100644 index 0000000000..9772fdf29c --- /dev/null +++ b/src/main/java/io/streamnative/kop/coordinator/group/package-info.java @@ -0,0 +1,17 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +/** + * Classes for kafka coordinator group. + */ +package io.streamnative.kop.coordinator.group; \ No newline at end of file diff --git a/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataTest.java b/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataTest.java new file mode 100644 index 0000000000..7b73e62f2a --- /dev/null +++ b/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataTest.java @@ -0,0 +1,390 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.coordinator.group; + +import static io.streamnative.kop.coordinator.group.GroupState.CompletingRebalance; +import static io.streamnative.kop.coordinator.group.GroupState.Dead; +import static io.streamnative.kop.coordinator.group.GroupState.Empty; +import static io.streamnative.kop.coordinator.group.GroupState.PreparingRebalance; +import static io.streamnative.kop.coordinator.group.GroupState.Stable; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertNull; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +import com.google.common.collect.Sets; +import java.util.LinkedHashMap; +import java.util.Map; +import java.util.Set; +import lombok.val; +import org.junit.Before; +import org.junit.Test; + +/** + * Unit test {@link GroupMetadata}. + */ +public class GroupMetadataTest { + + private static final String protocolType = "consumer"; + private static final String groupId = "test-group-id"; + private static final String clientId = "test-client-id"; + private static final String clientHost = "test-client-host"; + private static final int rebalanceTimeoutMs = 60000; + private static final int sessionTimeoutMs = 10000; + + private GroupMetadata group = null; + + @Before + public void setUp() { + group = new GroupMetadata(groupId, Empty); + } + + @Test + public void testCanRebalanceWhenStable() { + assertTrue(group.canReblance()); + } + + @Test + public void testCanRebalanceWhenCompletingRebalance() { + group.transitionTo(PreparingRebalance); + group.transitionTo(CompletingRebalance); + assertTrue(group.canReblance()); + } + + @Test + public void testCannotRebalanceWhenPreparingRebalance() { + group.transitionTo(PreparingRebalance); + assertFalse(group.canReblance()); + } + + @Test + public void testCannotRebalanceWhenDead() { + group.transitionTo(PreparingRebalance); + group.transitionTo(Empty); + group.transitionTo(Dead); + assertFalse(group.canReblance()); + } + + @Test + public void testStableToPreparingRebalanceTransition() { + group.transitionTo(PreparingRebalance); + assertState(group, PreparingRebalance); + } + + @Test + public void testStableToDeadTransition() { + group.transitionTo(Dead); + assertState(group, Dead); + } + + @Test + public void testAwaitingRebalanceToPreparingRebalanceTransition() { + group.transitionTo(PreparingRebalance); + group.transitionTo(CompletingRebalance); + group.transitionTo(PreparingRebalance); + assertState(group, PreparingRebalance); + } + + @Test + public void testPreparingRebalanceToDeadTransition() { + group.transitionTo(PreparingRebalance); + group.transitionTo(Dead); + assertState(group, Dead); + } + + @Test + public void testPreparingRebalanceToEmptyTransition() { + group.transitionTo(PreparingRebalance); + group.transitionTo(Empty); + assertState(group, Empty); + } + + @Test + public void testEmptyToDeadTransition() { + group.transitionTo(PreparingRebalance); + group.transitionTo(Empty); + group.transitionTo(Dead); + assertState(group, Dead); + } + + @Test + public void testAwaitingRebalanceToStableTransition() { + group.transitionTo(PreparingRebalance); + group.transitionTo(CompletingRebalance); + group.transitionTo(Stable); + assertState(group, Stable); + } + + @Test(expected = IllegalStateException.class) + public void testEmptyToStableIllegalTransition() { + group.transitionTo(Stable); + } + + @Test(expected = IllegalStateException.class) + public void testStableToStableIllegalTransition() { + group.transitionTo(PreparingRebalance); + group.transitionTo(CompletingRebalance); + group.transitionTo(Stable); + group.transitionTo(Stable); + fail("should have failed due to illegal transition"); + } + + @Test(expected = IllegalStateException.class) + public void testEmptyToAwaitingRebalanceIllegalTransition() { + group.transitionTo(CompletingRebalance); + } + + @Test(expected = IllegalStateException.class) + public void testPreparingRebalanceToPreparingRebalanceIllegalTransition() { + group.transitionTo(PreparingRebalance); + group.transitionTo(PreparingRebalance); + } + + @Test(expected = IllegalStateException.class) + public void testPreparingRebalanceToStableIllegalTransition() { + group.transitionTo(PreparingRebalance); + group.transitionTo(Stable); + } + + @Test(expected = IllegalStateException.class) + public void testAwaitingRebalanceToAwaitingRebalanceIllegalTransition() { + group.transitionTo(PreparingRebalance); + group.transitionTo(CompletingRebalance); + group.transitionTo(CompletingRebalance); + } + + @Test + public void testDeadToDeadIllegalTransition() { + group.transitionTo(PreparingRebalance); + group.transitionTo(Dead); + group.transitionTo(Dead); + assertState(group, Dead); + } + + @Test(expected = IllegalStateException.class) + public void testDeadToStableIllegalTransition() { + group.transitionTo(PreparingRebalance); + group.transitionTo(Dead); + group.transitionTo(Stable); + } + + @Test(expected = IllegalStateException.class) + public void testDeadToPreparingRebalanceIllegalTransition() { + group.transitionTo(PreparingRebalance); + group.transitionTo(Dead); + group.transitionTo(PreparingRebalance); + } + + @Test(expected = IllegalStateException.class) + public void testDeadToAwaitingRebalanceIllegalTransition() { + group.transitionTo(PreparingRebalance); + group.transitionTo(Dead); + group.transitionTo(CompletingRebalance); + } + + @Test + public void testSelectProtocol() { + String memberId = "memberId"; + Map protocols = new LinkedHashMap<>(); + protocols.put("range", new byte[0]); + protocols.put("roundrobin", new byte[0]); + MemberMetadata member = new MemberMetadata( + memberId, + groupId, + clientId, + clientHost, + rebalanceTimeoutMs, + sessionTimeoutMs, + protocolType, + protocols); + + group.add(member); + assertEquals("range", group.selectProtocol()); + + String otherMemberId = "otherMemberId"; + protocols = new LinkedHashMap<>(); + protocols.put("roundrobin", new byte[0]); + protocols.put("range", new byte[0]); + MemberMetadata otherMember = new MemberMetadata( + otherMemberId, + groupId, + clientId, + clientHost, + rebalanceTimeoutMs, + sessionTimeoutMs, + protocolType, + protocols); + + group.add(otherMember); + // now could be either range or robin since there is no majority preference + assertTrue(protocols.keySet().contains(group.selectProtocol())); + + String lastMemberId = "lastMemberId"; + protocols = new LinkedHashMap<>(); + protocols.put("roundrobin", new byte[0]); + protocols.put("range", new byte[0]); + val lastMember = new MemberMetadata( + lastMemberId, + groupId, + clientId, + clientHost, + rebalanceTimeoutMs, + sessionTimeoutMs, + protocolType, + protocols); + + group.add(lastMember); + // now we should prefer 'roundrobin' + assertEquals("roundrobin", group.selectProtocol()); + } + + @Test(expected = IllegalStateException.class) + public void testSelectProtocolRaisesIfNoMembers() { + group.selectProtocol(); + fail("Should not reach here"); + } + + @Test + public void testSelectProtocolChoosesCompatibleProtocol() { + String memberId = "memberId"; + Map protocols = new LinkedHashMap<>(); + protocols.put("range", new byte[0]); + protocols.put("roundrobin", new byte[0]); + MemberMetadata member = new MemberMetadata( + memberId, + groupId, + clientId, + clientHost, + rebalanceTimeoutMs, + sessionTimeoutMs, + protocolType, + protocols); + + String otherMemberId = "otherMemberId"; + protocols = new LinkedHashMap<>(); + protocols.put("roundrobin", new byte[0]); + protocols.put("blah", new byte[0]); + MemberMetadata otherMember = new MemberMetadata( + otherMemberId, + groupId, + clientId, + clientHost, + rebalanceTimeoutMs, + sessionTimeoutMs, + protocolType, + protocols); + + group.add(member); + group.add(otherMember); + assertEquals("roundrobin", group.selectProtocol()); + } + + @Test + public void testSupportsProtocols() { + // by default, the group supports everything + assertTrue(group.supportsProtocols(Sets.newHashSet("roundrobin", "range"))); + + String memberId = "memberId"; + Map protocols = new LinkedHashMap<>(); + protocols.put("range", new byte[0]); + protocols.put("roundrobin", new byte[0]); + MemberMetadata member = new MemberMetadata( + memberId, + groupId, + clientId, + clientHost, + rebalanceTimeoutMs, + sessionTimeoutMs, + protocolType, + protocols); + + group.add(member); + assertTrue(group.supportsProtocols(Sets.newHashSet("roundrobin", "foo"))); + assertTrue(group.supportsProtocols(Sets.newHashSet("range", "foo"))); + assertFalse(group.supportsProtocols(Sets.newHashSet("foo", "bar"))); + + String otherMemberId = "otherMemberId"; + protocols = new LinkedHashMap<>(); + protocols.put("roundrobin", new byte[0]); + protocols.put("blah", new byte[0]); + MemberMetadata otherMember = new MemberMetadata( + otherMemberId, + groupId, + clientId, + clientHost, + rebalanceTimeoutMs, + sessionTimeoutMs, + protocolType, + protocols); + + group.add(otherMember); + + assertTrue(group.supportsProtocols(Sets.newHashSet("roundrobin", "foo"))); + assertFalse(group.supportsProtocols(Sets.newHashSet("range", "foo"))); + } + + @Test + public void testInitNextGeneration() { + val memberId = "memberId"; + Map protocols = new LinkedHashMap<>(); + protocols.put("roundrobin", new byte[0]); + val member = new MemberMetadata( + memberId, + groupId, + clientId, + clientHost, + rebalanceTimeoutMs, + sessionTimeoutMs, + protocolType, + protocols); + + group.transitionTo(PreparingRebalance); + member.awaitingJoinCallback(e -> {}); + group.add(member); + + assertEquals(0, group.generationId()); + assertNull(group.protocolOrNull()); + + group.initNextGeneration(); + + assertEquals(1, group.generationId()); + assertEquals("roundrobin", group.protocolOrNull()); + } + + @Test + public void testInitNextGenerationEmptyGroup() { + assertEquals(Empty, group.currentState()); + assertEquals(0, group.generationId()); + assertNull(group.protocolOrNull()); + + group.transitionTo(PreparingRebalance); + group.initNextGeneration(); + + assertEquals(1, group.generationId()); + assertNull(group.protocolOrNull()); + } + + private void assertState(GroupMetadata group, GroupState targetState) { + Set states = Sets.newHashSet( + Stable, PreparingRebalance, CompletingRebalance, Dead + ); + Set otherStates = Sets.newHashSet(states); + otherStates.remove(targetState); + otherStates.forEach(otherState -> assertFalse(group.is(otherState))); + assertTrue(group.is(targetState)); + } + + +} diff --git a/src/test/java/io/streamnative/kop/coordinator/group/MemberMetadataTest.java b/src/test/java/io/streamnative/kop/coordinator/group/MemberMetadataTest.java new file mode 100644 index 0000000000..dbf79ff5fa --- /dev/null +++ b/src/test/java/io/streamnative/kop/coordinator/group/MemberMetadataTest.java @@ -0,0 +1,126 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.coordinator.group; + +import static org.junit.Assert.assertArrayEquals; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; +import static org.junit.Assert.fail; + +import com.google.common.collect.Sets; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.Map; +import org.junit.Test; +import org.testng.collections.Maps; + +/** + * Unit test of {@link MemberMetadata}. + */ +public class MemberMetadataTest { + + private static final String groupId = "test-group-id"; + private static final String clientId = "test-client-id"; + private static final String clientHost = "test-client-host"; + private static final String memberId = "test-member-id"; + private static final String protocolType = "consumer"; + private static final int rebalanceTimeoutMs = 60000; + private static final int sessionTimemoutMs = 10000; + + private static MemberMetadata newMember(Map protocols) { + return new MemberMetadata( + memberId, + groupId, + clientId, + clientHost, + rebalanceTimeoutMs, + sessionTimemoutMs, + protocolType, + protocols + ); + } + + @Test + public void testMatchesSupportedProtocols() { + Map protocols = Maps.newHashMap(); + protocols.put("range", new byte[0]); + + MemberMetadata member = newMember(protocols); + + assertTrue(member.matches(protocols)); + + protocols = new HashMap<>(); + protocols.put("range", new byte[] { 0 }); + assertFalse(member.matches(protocols)); + + protocols = new HashMap<>(); + protocols.put("roundrobin", new byte[0]); + assertFalse(member.matches(protocols)); + + protocols = new HashMap<>(); + protocols.put("range", new byte[0]); + protocols.put("roundrobin", new byte[0]); + assertFalse(member.matches(protocols)); + } + + @Test + public void testVoteForPreferredProtocol() { + Map protocols = new LinkedHashMap<>(); + protocols.put("range", new byte[0]); + protocols.put("roundrobin", new byte[0]); + + MemberMetadata member = newMember(protocols); + assertEquals("range", member.vote(Sets.newHashSet("range", "roundrobin"))); + assertEquals("roundrobin", member.vote(Sets.newHashSet("blah", "roundrobin"))); + } + + @Test + public void testMetadata() { + Map protocols = new LinkedHashMap<>(); + protocols.put("range", new byte[] { 0xf }); + protocols.put("roundrobin", new byte[] { 0xe }); + + MemberMetadata memberMetadata = newMember(protocols); + assertArrayEquals( + new byte[] { 0xf }, + memberMetadata.metadata("range")); + assertArrayEquals( + new byte[] { 0xe }, + memberMetadata.metadata("roundrobin")); + } + + @Test(expected = IllegalArgumentException.class) + public void testMetadataRaisesOnUnsupportedProtocol() { + Map protocols = new LinkedHashMap<>(); + protocols.put("range", new byte[0]); + protocols.put("roundrobin", new byte[0]); + + MemberMetadata member = newMember(protocols); + member.metadata("blah"); + fail("Should not reach here"); + } + + @Test(expected = IllegalArgumentException.class) + public void testVoteRaisesOnUnsupportedProtocols() { + Map protocols = new LinkedHashMap<>(); + protocols.put("range", new byte[0]); + protocols.put("roundrobin", new byte[0]); + + MemberMetadata member = newMember(protocols); + member.vote(Sets.newHashSet("blah")); + fail("Should not reach here"); + } + +} From 718096aa606cc98a340bd5886059ee72692ba2ae Mon Sep 17 00:00:00 2001 From: Sijie Guo Date: Tue, 23 Jul 2019 22:51:15 +0800 Subject: [PATCH 2/8] Group coordindator part 2 --- .../kop/coordinator/group/GroupMetadata.java | 18 ++ .../group/GroupMetadataConstants.java | 265 ++++++++++++++++++ .../group/GroupMetadataManager.java | 206 ++++++++++++++ .../kop/coordinator/group/OffsetConfig.java | 6 + .../io/streamnative/kop/utils/CoreUtils.java | 28 ++ 5 files changed, 523 insertions(+) create mode 100644 src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataConstants.java create mode 100644 src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataManager.java create mode 100644 src/main/java/io/streamnative/kop/utils/CoreUtils.java diff --git a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadata.java b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadata.java index aca6780093..021c5ea274 100644 --- a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadata.java +++ b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadata.java @@ -18,8 +18,10 @@ import com.google.common.base.MoreObjects; import com.google.common.base.MoreObjects.ToStringHelper; +import com.google.common.base.Supplier; import com.google.common.collect.Sets; import io.streamnative.kop.coordinator.group.MemberMetadata.MemberSummary; +import io.streamnative.kop.utils.CoreUtils; import java.util.Collections; import java.util.Comparator; import java.util.HashMap; @@ -158,14 +160,30 @@ static class GroupSummary { this.state = initialState; } + public T inLock(Supplier supplier) { + return CoreUtils.inLock(lock, supplier); + } + + public Optional protocolType() { + return protocolType; + } + public GroupState currentState() { return state; } + public String groupId() { + return groupId; + } + public long generationId() { return generationId; } + public List allMemberMetadata() { + return members.values().stream().collect(Collectors.toList()); + } + public boolean is(GroupState groupState) { return state == groupState; } diff --git a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataConstants.java b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataConstants.java new file mode 100644 index 0000000000..fa6b06ebff --- /dev/null +++ b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataConstants.java @@ -0,0 +1,265 @@ +package io.streamnative.kop.coordinator.group; + +import static com.google.common.base.Preconditions.checkState; +import static org.apache.kafka.common.protocol.types.Type.BYTES; +import static org.apache.kafka.common.protocol.types.Type.INT32; +import static org.apache.kafka.common.protocol.types.Type.INT64; +import static org.apache.kafka.common.protocol.types.Type.NULLABLE_STRING; +import static org.apache.kafka.common.protocol.types.Type.STRING; + +import com.google.common.collect.Lists; +import java.nio.ByteBuffer; +import java.util.List; +import java.util.Map; +import java.util.stream.Collectors; +import lombok.val; +import org.apache.kafka.common.KafkaException; +import org.apache.kafka.common.protocol.types.ArrayOf; +import org.apache.kafka.common.protocol.types.BoundField; +import org.apache.kafka.common.protocol.types.Field; +import org.apache.kafka.common.protocol.types.Schema; +import org.apache.kafka.common.protocol.types.Struct; +import org.apache.pulsar.common.schema.KeyValue; + +/** + * Messages stored for the group topic has versions for both the key and value fields. Key + * version is used to indicate the type of the message (also to differentiate different types + * of messages from being compacted together if they have the same field values); and value + * version is used to evolve the messages within their data types: + * + *

key version 0: group consumption offset + * -> value version 0: [offset, metadata, timestamp] + * + *

key version 1: group consumption offset + * -> value version 1: [offset, metadata, commit_timestamp, expire_timestamp] + * + *

key version 2: group metadata + * -> value version 0: [protocol_type, generation, protocol, leader, members] + */ +final class GroupMetadataConstants { + + static final short CURRENT_OFFSET_KEY_SCHEMA_VERSION = 1; + static final short CURRENT_GROUP_KEY_SCHEMA_VERSION = 2; + + static final Schema OFFSET_COMMIT_KEY_SCHEMA = new Schema( + new Field("group", STRING), + new Field("topic", STRING), + new Field("partition", INT32) + ); + static final BoundField OFFSET_KEY_GROUP_FIELD = OFFSET_COMMIT_KEY_SCHEMA.get("group"); + static final BoundField OFFSET_KEY_TOPIC_FIELD = OFFSET_COMMIT_KEY_SCHEMA.get("topic"); + static final BoundField OFFSET_KEY_PARTITION_FIELD = OFFSET_COMMIT_KEY_SCHEMA.get("partition"); + + static final Schema OFFSET_COMMIT_VALUE_SCHEMA_V0 = new Schema( + new Field("offset", INT64), + new Field("metadata", STRING, "Associated metadata.", ""), + new Field("timestamp", INT64) + ); + static final BoundField OFFSET_VALUE_OFFSET_FIELD_V0 = OFFSET_COMMIT_VALUE_SCHEMA_V0.get("offset"); + static final BoundField OFFSET_VALUE_METADATA_FIELD_V0 = OFFSET_COMMIT_VALUE_SCHEMA_V0.get("metadata"); + static final BoundField OFFSET_VALUE_TIMESTAMP_FIELD_V0 = OFFSET_COMMIT_VALUE_SCHEMA_V0.get("timestamp"); + + static final Schema OFFSET_COMMIT_VALUE_SCHEMA_V1 = new Schema( + new Field("offset", INT64), + new Field("metadata", STRING, "Associated metadata.", ""), + new Field("commit_timestamp", INT64), + new Field("expire_timestamp", INT64) + ); + static final BoundField OFFSET_VALUE_OFFSET_FIELD_V1 = OFFSET_COMMIT_VALUE_SCHEMA_V1.get("offset"); + static final BoundField OFFSET_VALUE_METADATA_FIELD_V1 = OFFSET_COMMIT_VALUE_SCHEMA_V1.get("metadata"); + static final BoundField OFFSET_VALUE_COMMIT_TIMESTAMP_FIELD_V1 = + OFFSET_COMMIT_VALUE_SCHEMA_V1.get("commit_timestamp"); + static final BoundField OFFSET_VALUE_EXPIRE_TIMESTAMP_FIELD_V1 = + OFFSET_COMMIT_VALUE_SCHEMA_V1.get("expire_timestamp"); + + static final Schema GROUP_METADATA_KEY_SCHEMA = new Schema(new Field("group", STRING)); + static final BoundField GROUP_KEY_GROUP_FIELD = GROUP_METADATA_KEY_SCHEMA.get("group"); + + static final String MEMBER_ID_KEY = "member_id"; + static final String CLIENT_ID_KEY = "client_id"; + static final String CLIENT_HOST_KEY = "client_host"; + static final String REBALANCE_TIMEOUT_KEY = "rebalance_timeout"; + static final String SESSION_TIMEOUT_KEY = "session_timeout"; + static final String SUBSCRIPTION_KEY = "subscription"; + static final String ASSIGNMENT_KEY = "assignment"; + + static final Schema MEMBER_METADATA_V0 = new Schema( + new Field(MEMBER_ID_KEY, STRING), + new Field(CLIENT_ID_KEY, STRING), + new Field(CLIENT_HOST_KEY, STRING), + new Field(SESSION_TIMEOUT_KEY, INT32), + new Field(SUBSCRIPTION_KEY, BYTES), + new Field(ASSIGNMENT_KEY, BYTES)); + + static final Schema MEMBER_METADATA_V1 = new Schema( + new Field(MEMBER_ID_KEY, STRING), + new Field(CLIENT_ID_KEY, STRING), + new Field(CLIENT_HOST_KEY, STRING), + new Field(REBALANCE_TIMEOUT_KEY, INT32), + new Field(SESSION_TIMEOUT_KEY, INT32), + new Field(SUBSCRIPTION_KEY, BYTES), + new Field(ASSIGNMENT_KEY, BYTES)); + + static final String PROTOCOL_TYPE_KEY = "protocol_type"; + static final String GENERATION_KEY = "generation"; + static final String PROTOCOL_KEY = "protocol"; + static final String LEADER_KEY = "leader"; + static final String MEMBERS_KEY = "members"; + + static final Schema GROUP_METADATA_VALUE_SCHEMA_V0 = new Schema( + new Field(PROTOCOL_TYPE_KEY, STRING), + new Field(GENERATION_KEY, INT32), + new Field(PROTOCOL_KEY, NULLABLE_STRING), + new Field(LEADER_KEY, NULLABLE_STRING), + new Field(MEMBERS_KEY, new ArrayOf(MEMBER_METADATA_V0))); + + static final Schema GROUP_METADATA_VALUE_SCHEMA_V1 = new Schema( + new Field(PROTOCOL_TYPE_KEY, STRING), + new Field(GENERATION_KEY, INT32), + new Field(PROTOCOL_KEY, NULLABLE_STRING), + new Field(LEADER_KEY, NULLABLE_STRING), + new Field(MEMBERS_KEY, new ArrayOf(MEMBER_METADATA_V1))); + + // map of versions to key schemas as data types + static final Map MESSAGE_TYPE_SCHEMAS = asMap( + kv(0, OFFSET_COMMIT_KEY_SCHEMA), + kv(1, OFFSET_COMMIT_KEY_SCHEMA), + kv(2, GROUP_METADATA_KEY_SCHEMA) + ); + + // map of version of offset value schemas + static final Map OFFSET_VALUE_SCHEMAS = asMap( + kv(0, OFFSET_COMMIT_VALUE_SCHEMA_V0), + kv(1, OFFSET_COMMIT_VALUE_SCHEMA_V1) + ); + static final short CURRENT_OFFSET_VALUE_SCHEMA_VERSION = 1; + + // map of version of group metadata value schemas + static final Map GROUP_VALUE_SCHEMAS = asMap( + kv(0, GROUP_METADATA_VALUE_SCHEMA_V0), + kv(1, GROUP_METADATA_VALUE_SCHEMA_V1) + ); + static final short CURRENT_GROUP_VALUE_SCHEMA_VERSION = 1; + + static final Schema CURRENT_OFFSET_KEY_SCHEMA = schemaForKey(CURRENT_OFFSET_KEY_SCHEMA_VERSION); + static final Schema CURRENT_GROUP_KEY_SCHEMA = schemaForKey(CURRENT_GROUP_KEY_SCHEMA_VERSION); + + static final Schema CURRENT_OFFSET_VALUE_SCHEMA = schemaForOffset(CURRENT_OFFSET_VALUE_SCHEMA_VERSION); + static final Schema CURRENT_GROUP_VALUE_SCHEMA = schemaForGroup(CURRENT_GROUP_VALUE_SCHEMA_VERSION); + + private static final Schema schemaForKey(int version) { + Schema schema = MESSAGE_TYPE_SCHEMAS.get(version); + if (null == schema) { + throw new KafkaException("Unknown offset schema version " + version); + } + return schema; + } + + private static final Schema schemaForOffset(int version) { + Schema schema = OFFSET_VALUE_SCHEMAS.get(version); + if (null == schema) { + throw new KafkaException("Unknown offset schema version " + version); + } + return schema; + } + + private static final Schema schemaForGroup(int version) { + Schema schema = GROUP_VALUE_SCHEMAS.get(version); + if (null == schema) { + throw new KafkaException("Unknown group metadata version " + version); + } + return schema; + } + + private static KeyValue kv(K key, V value) { + return new KeyValue<>(key, value); + } + + private static Map asMap(KeyValue ...kvs) { + return Lists.newArrayList(kvs) + .stream() + .collect(Collectors.toMap( + e -> e.getKey(), + e -> e.getValue() + )); + } + + /** + * Generates the key for group metadata message for given group + * + * @return key bytes for group metadata message + */ + static byte[] groupMetadataKey(String group) { + Struct key = new Struct(CURRENT_GROUP_KEY_SCHEMA); + key.set(GROUP_KEY_GROUP_FIELD, group); + ByteBuffer byteBuffer = ByteBuffer.allocate(2 /* version */ + key.sizeOf()); + byteBuffer.putShort(CURRENT_GROUP_KEY_SCHEMA_VERSION); + key.writeTo(byteBuffer); + return byteBuffer.array(); + } + + /** + * Generates the payload for group metadata message from given offset and metadata + * assuming the generation id, selected protocol, leader and member assignment are all available + * + * @param groupMetadata current group metadata + * @param assignment the assignment for the rebalancing generation + * @param version the version of the value message to use + * @return payload for offset commit message + */ + static byte[] groupMetadataValue(GroupMetadata groupMetadata, + Map assignment, + short version) { + Struct value; + if (version == 0) { + value = new Struct(GROUP_METADATA_VALUE_SCHEMA_V0); + } else { + value = new Struct(CURRENT_GROUP_VALUE_SCHEMA); + } + + value.set(PROTOCOL_TYPE_KEY, groupMetadata.protocolType().orElse("")); + value.set(GENERATION_KEY, groupMetadata.generationId()); + value.set(PROTOCOL_KEY, groupMetadata.protocolOrNull()); + value.set(LEADER_KEY, groupMetadata.leaderOrNull()); + + List memberStructs = groupMetadata.allMemberMetadata().stream().map(memberMetadata -> { + Struct memberStruct = value.instance(MEMBERS_KEY); + memberStruct.set(MEMBER_ID_KEY, memberMetadata.memberId()); + memberStruct.set(CLIENT_ID_KEY, memberMetadata.clientId()); + memberStruct.set(CLIENT_HOST_KEY, memberMetadata.clientHost()); + memberStruct.set(SESSION_TIMEOUT_KEY, memberMetadata.sessionTimeoutMs()); + + if (version > 0) { + memberStruct.set(REBALANCE_TIMEOUT_KEY, memberMetadata.rebalanceTimeoutMs()); + } + + // The group is non-empty, so the current protocol must be defined + String protocol = groupMetadata.protocolOrNull(); + if (protocol == null) { + throw new IllegalStateException("Attempted to write non-empty group metadata with no defined protocol"); + } + + byte[] metadata = memberMetadata.metadata(protocol); + memberStruct.set(SUBSCRIPTION_KEY, ByteBuffer.wrap(metadata)); + + byte[] memberAssignment = assignment.get(memberMetadata.memberId()); + checkState( + memberAssignment != null, + "Member assignment is null for member %s", memberMetadata.memberId()); + + memberStruct.set(ASSIGNMENT_KEY, ByteBuffer.wrap(memberAssignment)); + + return memberStruct; + }).collect(Collectors.toList()); + + value.set(MEMBERS_KEY, memberStructs.toArray()); + + ByteBuffer byteBuffer = ByteBuffer.allocate(2 /* version */ + value.sizeOf()); + byteBuffer.putShort(version); + value.writeTo(byteBuffer); + return byteBuffer.array(); + } + + private GroupMetadataConstants() {} + +} diff --git a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataManager.java b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataManager.java new file mode 100644 index 0000000000..5d72e725dd --- /dev/null +++ b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataManager.java @@ -0,0 +1,206 @@ +package io.streamnative.kop.coordinator.group; + +import static io.streamnative.kop.coordinator.group.GroupMetadataConstants.CURRENT_GROUP_KEY_SCHEMA; +import static io.streamnative.kop.coordinator.group.GroupMetadataConstants.CURRENT_GROUP_KEY_SCHEMA_VERSION; +import static io.streamnative.kop.coordinator.group.GroupMetadataConstants.GROUP_KEY_GROUP_FIELD; +import static io.streamnative.kop.utils.CoreUtils.inLock; +import static java.nio.charset.StandardCharsets.UTF_8; + +import java.nio.ByteBuffer; +import java.util.HashSet; +import java.util.Map; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.locks.ReentrantLock; +import java.util.function.Consumer; +import lombok.Data; +import lombok.experimental.Accessors; +import org.apache.bookkeeper.common.hash.Murmur3; +import org.apache.bookkeeper.common.util.MathUtils; +import org.apache.bookkeeper.common.util.OrderedScheduler; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.protocol.types.Struct; +import org.apache.kafka.common.record.CompressionType; +import org.apache.kafka.common.requests.ApiVersionsResponse.ApiVersion; + +/** + * Manager to manage a coordination group. + */ +class GroupMetadataManager { + + /** + * The key interface. + */ + interface BaseKey { + short version(); + Object key(); + } + + /** + * Key to index group metadata. + */ + @Data + @Accessors(fluent = true) + static class GroupMetadataKey { + + private final short version; + private final String key; + + @Override + public String toString() { + return key; + } + + } + + /** + * The group on a topic partition. + */ + @Data + @Accessors(fluent = true) + static class GroupTopicPartition { + + private final String group; + private final TopicPartition topicPartition; + + GroupTopicPartition(String group, + String topic, + int partition) { + this.group = group; + this.topicPartition = new TopicPartition(topic, partition); + } + + @Override + public String toString() { + return String.format( + "[%s, %s, %d]", + group, topicPartition.topic(), topicPartition.partition() + ); + } + + } + + private final int brokerId; + private final ApiVersion interBrokerProtocolVersion; + private final OffsetConfig config; + private final CompressionType compressionType; + private final ConcurrentMap groupMetadataCache; + /* lock protecting access to loading and owned partition sets */ + private final ReentrantLock partitionLock = new ReentrantLock(); + /** + * partitions of consumer groups that are being loaded, its lock should + * be always called BEFORE the group lock if needed + */ + private final Set loadingPartitions = new HashSet<>(); + /* partitions of consumer groups that are assigned, using the same loading partition lock */ + private final Set ownedPartitions = new HashSet<>(); + /* shutting down flag */ + private final AtomicBoolean shuttingDown = new AtomicBoolean(false); + /* single-thread scheduler to handle offset/group metadata cache loading and unloading */ + private OrderedScheduler scheduler; + private final String logIdent; + private final int groupMetadataTopicPartitionCount; + + GroupMetadataManager(int brokerId, + ApiVersion interBrokerProtocolVersion, + OffsetConfig config, + int groupMetadataTopicPartitionCount) { + this.brokerId = brokerId; + this.interBrokerProtocolVersion = interBrokerProtocolVersion; + this.config = config; + this.compressionType = config.offsetsTopicCompressionType(); + this.groupMetadataCache = new ConcurrentHashMap<>(); + this.logIdent = String.format("[GroupMetadataManager brokerId=%d]", brokerId); + this.groupMetadataTopicPartitionCount = groupMetadataTopicPartitionCount; + } + + public void startup() { + this.scheduler = OrderedScheduler.newSchedulerBuilder() + .name("group-metadata-manager") + .numThreads(1) + .build(); + } + + public Iterable currentGroups() { + return groupMetadataCache.values(); + } + + + public boolean isPartitionOwned(int partition) { + return inLock( + partitionLock, + () -> ownedPartitions.contains(partition)); + } + + public boolean isPartitionLoading(int partition) { + return inLock( + partitionLock, + () -> loadingPartitions.contains(partition) + ); + } + + public int partitionFor(String groupId) { + return MathUtils.signSafeMod( + Murmur3.hash32(groupId.getBytes(UTF_8)), + groupMetadataTopicPartitionCount + ); + } + + public boolean isGroupLocal(String groupId) { + return isPartitionOwned(partitionFor(groupId)); + } + + public boolean isGroupLoading(String groupId) { + return isPartitionLoading(partitionFor(groupId)); + } + + public boolean isLoading() { + return inLock( + partitionLock, + () -> !loadingPartitions.isEmpty() + ); + } + + // return true iff group is owned and the group doesn't exist + public boolean groupNotExists(String groupId) { + return inLock( + partitionLock, + () -> { + if (isGroupLocal(groupId)) { + return true; + } else { + return getGroup(groupId) + .map(metadata -> metadata.inLock(() -> metadata.is(GroupState.Dead))) + .orElse(false); + } + } + ); + } + + public Optional getGroup(String groupId) { + return Optional.ofNullable(groupMetadataCache.getOrDefault(groupId, null)); + } + + public GroupMetadata addGroup(GroupMetadata group) { + GroupMetadata oldGroup = groupMetadataCache.putIfAbsent(group.groupId(), group); + if (null != oldGroup) { + return oldGroup; + } else { + return group; + } + } + + public CompletableFuture storeGroup(GroupMetadata group, + Map groupAssignment) { + + } + + + + +} diff --git a/src/main/java/io/streamnative/kop/coordinator/group/OffsetConfig.java b/src/main/java/io/streamnative/kop/coordinator/group/OffsetConfig.java index c90f02b64c..cbf172480f 100644 --- a/src/main/java/io/streamnative/kop/coordinator/group/OffsetConfig.java +++ b/src/main/java/io/streamnative/kop/coordinator/group/OffsetConfig.java @@ -15,11 +15,17 @@ import lombok.Builder; import lombok.Data; +import lombok.experimental.Accessors; +import org.apache.kafka.common.record.CompressionType; /** * Offset configuration. */ @Builder @Data +@Accessors(fluent = true) public class OffsetConfig { + + private final CompressionType offsetsTopicCompressionType; + } diff --git a/src/main/java/io/streamnative/kop/utils/CoreUtils.java b/src/main/java/io/streamnative/kop/utils/CoreUtils.java new file mode 100644 index 0000000000..1ed7afbeed --- /dev/null +++ b/src/main/java/io/streamnative/kop/utils/CoreUtils.java @@ -0,0 +1,28 @@ +package io.streamnative.kop.utils; + +import java.util.concurrent.locks.Lock; +import java.util.function.Supplier; + +/** + * Core utility functions. + */ +public final class CoreUtils { + + /** + * Retrieve a value under the protection of a lock. + * + * @param lock the lock to protect the operation + * @param supplier the supplier function to return the value + * @return the value retrieved from the function + */ + public static T inLock(Lock lock, Supplier supplier) { + lock.lock(); + try { + return supplier.get(); + } finally { + lock.unlock(); + } + } + + private CoreUtils() {} +} From 61e7b41dcafe8e75196f132de247703ee2dac38d Mon Sep 17 00:00:00 2001 From: Sijie Guo Date: Sat, 3 Aug 2019 13:14:44 +0800 Subject: [PATCH 3/8] Implement Group Coordinator --- pom.xml | 59 +- .../coordinator/group/GroupCoordinator.java | 842 +++++++++++++++++- .../kop/coordinator/group/GroupMetadata.java | 45 +- .../group/GroupMetadataConstants.java | 94 ++ .../group/GroupMetadataManager.java | 231 ++++- .../kop/coordinator/group/MemberMetadata.java | 8 +- .../group/GroupMetadataManagerTest.java | 10 + .../coordinator/group/GroupMetadataTest.java | 3 +- 8 files changed, 1264 insertions(+), 28 deletions(-) create mode 100644 src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataManagerTest.java diff --git a/pom.xml b/pom.xml index 7df9da7152..0e8f1e7dd9 100644 --- a/pom.xml +++ b/pom.xml @@ -39,7 +39,9 @@ 1.18.4 2.22.0 4.1.32.Final - 2.4.0 + 2.5.0-ef274bbb8 + 2.11 + 2.11.12 1.7.25 3.1.8 1.11.2 @@ -52,6 +54,7 @@ 3.0.0-M1 1.4.1.Final 6.19 + 4.0.2 3.1.8 @@ -93,6 +96,12 @@ + + org.scala-lang + scala-library + ${scala.version} + + org.apache.pulsar pulsar-websocket @@ -307,6 +316,54 @@ + + + net.alchim31.maven + scala-maven-plugin + ${scala-maven-plugin.version} + + + eclipse-add-source + + add-source + + + + scala-compile-first + + compile + + + + scala-test-compile-first + + testCompile + + + + + ${scala.binary.version} + all + + -unchecked + -deprecation + -feature + -explaintypes + -Yno-adapted-args + + + -Xms1024m + -Xmx1024m + + + -source + ${javac.target} + -target + ${javac.target} + -Xlint:all,-serial,-path,-try + + + diff --git a/src/main/java/io/streamnative/kop/coordinator/group/GroupCoordinator.java b/src/main/java/io/streamnative/kop/coordinator/group/GroupCoordinator.java index ab631869b6..adc9f65fc5 100644 --- a/src/main/java/io/streamnative/kop/coordinator/group/GroupCoordinator.java +++ b/src/main/java/io/streamnative/kop/coordinator/group/GroupCoordinator.java @@ -13,12 +13,46 @@ */ package io.streamnative.kop.coordinator.group; +import static com.google.common.base.Preconditions.checkState; +import static io.streamnative.kop.coordinator.group.GroupState.CompletingRebalance; +import static io.streamnative.kop.coordinator.group.GroupState.Dead; +import static io.streamnative.kop.coordinator.group.GroupState.Empty; +import static io.streamnative.kop.coordinator.group.GroupState.PreparingRebalance; +import static io.streamnative.kop.coordinator.group.GroupState.Stable; + +import com.google.common.collect.Sets; +import io.streamnative.kop.coordinator.group.GroupMetadata.GroupOverview; import io.streamnative.kop.coordinator.group.GroupMetadata.GroupSummary; +import io.streamnative.kop.utils.CoreUtils; +import java.util.ArrayList; import java.util.Collections; +import java.util.HashMap; +import java.util.HashSet; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.Set; +import java.util.concurrent.CompletableFuture; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.function.BiConsumer; +import java.util.function.Supplier; +import java.util.stream.Collectors; +import lombok.extern.slf4j.Slf4j; +import lombok.val; +import org.apache.kafka.common.internals.Topic; +import org.apache.kafka.common.protocol.ApiKeys; +import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.requests.JoinGroupRequest; +import org.apache.kafka.common.utils.Time; +import org.apache.pulsar.common.schema.KeyValue; +import org.apache.pulsar.common.util.FutureUtil; +import scala.Array; /** * Group coordinator. */ +@Slf4j public class GroupCoordinator { static final String NoState = ""; @@ -34,11 +68,817 @@ public class GroupCoordinator { Collections.emptyList() ); static final GroupSummary DeadGroup = new GroupSummary( - GroupState.Dead.toString(), + Dead.toString(), NoProtocolType, NoProtocol, Collections.emptyList() ); + private static boolean isValidGroupId(String groupId, + ApiKeys api) { + switch (api) { + case OFFSET_COMMIT: + case OFFSET_FETCH: + case DESCRIBE_GROUPS: + case DELETE_GROUPS: + return null != groupId; + default: + return null != groupId && !groupId.isEmpty(); + } + } + + private final AtomicBoolean isActive = new AtomicBoolean(false); + private final GroupConfig groupConfig; + private final GroupMetadataManager groupManager; + private final Time time; + + public CompletableFuture handleJoinGroup( + String groupId, + String memberId, + String clientId, + String clientHost, + int rebalanceTimeoutMs, + int sessionTimeoutMs, + String protocolType, + Map protocols + ) { + Optional errors = validateGroupStatus(groupId, ApiKeys.JOIN_GROUP); + if (errors.isPresent()) { + return CompletableFuture.completedFuture( + joinError(memberId, errors.get())); + } + + if (sessionTimeoutMs < groupConfig.getGroupMinSessionTimeoutMs() + || sessionTimeoutMs > groupConfig.getGroupMaxSessionTimeoutMs()) { + return CompletableFuture.completedFuture( + joinError(memberId, Errors.INVALID_SESSION_TIMEOUT)); + } else { + return groupManager.getGroup(groupId).map(group -> doJoinGroup( + group, + memberId, + clientId, + clientHost, + rebalanceTimeoutMs, + sessionTimeoutMs, + protocolType, + protocols + )).orElseGet(() -> { + if (memberId != JoinGroupRequest.UNKNOWN_MEMBER_ID) { + return CompletableFuture.completedFuture( + joinError(memberId, Errors.UNKNOWN_MEMBER_ID)); + } else { + GroupMetadata group = groupManager.addGroup(new GroupMetadata( + groupId, Empty + )); + return doJoinGroup( + group, + memberId, + clientId, + clientHost, + rebalanceTimeoutMs, + sessionTimeoutMs, + protocolType, + protocols + ); + } + }); + } + } + + private CompletableFuture doJoinGroup( + GroupMetadata group, + String memberId, + String clientId, + String clientHost, + int rebalanceTimeoutMs, + int sessionTimeoutMs, + String protocolType, + Map protocols + ) { + return group.inLock(() -> unsafeJoinGroup( + group, + memberId, + clientId, + clientHost, + rebalanceTimeoutMs, + sessionTimeoutMs, + protocolType, + protocols + )); + } + + private CompletableFuture unsafeJoinGroup( + GroupMetadata group, + String memberId, + String clientId, + String clientHost, + int rebalanceTimeoutMs, + int sessionTimeoutMs, + String protocolType, + Map protocols + ) { + if (!group.is(Empty) && ( + !group.protocolType().isPresent() + || group.protocolType().get() != protocolType + || !group.supportsProtocols(protocols.keySet()))) { + // if the new member does not support the group protocol, reject it + return CompletableFuture.completedFuture( + joinError(memberId, Errors.INCONSISTENT_GROUP_PROTOCOL)); + } else if (group.is(Empty) + && (protocols.isEmpty() || protocolType.isEmpty())) { + //reject if first member with empty group protocol or protocolType is empty + return CompletableFuture.completedFuture( + joinError(memberId, Errors.INCONSISTENT_GROUP_PROTOCOL)); + } else if (memberId != JoinGroupRequest.UNKNOWN_MEMBER_ID && !group.has(memberId)) { + // if the member trying to register with a un-recognized id, send the response to let + // it reset its member id and retry + return CompletableFuture.completedFuture( + joinError(memberId, Errors.UNKNOWN_MEMBER_ID)); + } else { + CompletableFuture resultFuture; + switch (group.currentState()) { + case Dead: + // if the group is marked as dead, it means some other thread has just removed the group + // from the coordinator metadata; this is likely that the group has migrated to some other + // coordinator OR the group is in a transient unstable phase. Let the member retry + // joining without the specified member id, + resultFuture = CompletableFuture.completedFuture( + joinError(memberId, Errors.UNKNOWN_MEMBER_ID)); + break; + case PreparingRebalance: + if (memberId == JoinGroupRequest.UNKNOWN_MEMBER_ID) { + resultFuture = addMemberAndRebalance( + rebalanceTimeoutMs, + sessionTimeoutMs, + clientId, + clientHost, + protocolType, + protocols, + group + ); + } else { + MemberMetadata member = group.get(memberId); + resultFuture = updateMemberAndRebalance(group, member, protocols); + } + break; + case CompletingRebalance: + if (JoinGroupRequest.UNKNOWN_MEMBER_ID == memberId) { + resultFuture = addMemberAndRebalance( + rebalanceTimeoutMs, + sessionTimeoutMs, + clientId, + clientHost, + protocolType, + protocols, + group + ); + } else { + MemberMetadata member = group.get(memberId); + if (member.matches(protocols)) { + // member is joining with the same metadata (which could be because it failed to + // receive the initial JoinGroup response), so just return current group information + // for the current generation. + Map members; + if (group.isLeader(memberId)) { + members = group.currentMemberMetadata(); + } else { + members = Collections.emptyMap(); + } + resultFuture = CompletableFuture.completedFuture( + new JoinGroupResult( + members, + memberId, + (int) group.generationId(), + group.protocolOrNull(), + group.leaderOrNull(), + Errors.NONE + ) + ); + } else { + resultFuture = updateMemberAndRebalance( + group, + member, + protocols + ); + } + } + break; + case Empty: + case Stable: + if (JoinGroupRequest.UNKNOWN_MEMBER_ID == memberId) { + // if the member id is unknown, register the member to the group + resultFuture = addMemberAndRebalance( + rebalanceTimeoutMs, + sessionTimeoutMs, + clientId, + clientHost, + protocolType, + protocols, + group + ); + } else { + MemberMetadata member = group.get(memberId); + if (group.isLeader(memberId) || !member.matches(protocols)) { + // force a rebalance if a member has changed metadata or if the leader sends JoinGroup. + // The latter allows the leader to trigger rebalances for changes affecting assignment + // which do not affect the member metadata (such as topic metadata changes for the consumer) + resultFuture = updateMemberAndRebalance(group, member, protocols); + } else { + // for followers with no actual change to their metadata, just return group information + // for the current generation which will allow them to issue SyncGroup + resultFuture = CompletableFuture.completedFuture(new JoinGroupResult( + Collections.emptyMap(), + memberId, + (int) group.generationId(), + group.protocolOrNull(), + group.leaderOrNull(), + Errors.NONE)); + } + } + break; + default: + resultFuture = FutureUtil.failedFuture( + new IllegalStateException("Unknown state " + group.currentState())); + break; + } + if (group.is(PreparingRebalance)) { + // TODO: check and trigger rebalance + } + return resultFuture; + } + + } + + public void handleSyncGroup(String groupId, + int generation, + String memberId, + Map groupAssignment, + BiConsumer responseCallback) { + Optional errorsOpt = validateGroupStatus(groupId, ApiKeys.SYNC_GROUP); + if (errorsOpt.isPresent()) { + Errors error = errorsOpt.get(); + if (Errors.COORDINATOR_LOAD_IN_PROGRESS == error) { + // The coordinator is loading, which means we've lost the state of the active rebalance and the + // group will need to start over at JoinGroup. By returning rebalance in progress, the consumer + // will attempt to rejoin without needing to rediscover the coordinator. Note that we cannot + // return COORDINATOR_LOAD_IN_PROGRESS since older clients do not expect the error. + responseCallback.accept(new byte[0], Errors.REBALANCE_IN_PROGRESS); + } else { + responseCallback.accept(new byte[0], error); + } + } else { + Optional groupOpt = groupManager.getGroup(groupId); + if (groupOpt.isPresent()) { + doSyncGroup( + groupOpt.get(), + generation, + memberId, + groupAssignment, + responseCallback + ); + } else { + responseCallback.accept(new byte[0], Errors.UNKNOWN_MEMBER_ID); + } + } + } + + private void doSyncGroup(GroupMetadata group, + int generationId, + String memberId, + final Map groupAssignment, + BiConsumer responseCallback) { + group.inLock(() -> { + if (!group.has(memberId)) { + responseCallback.accept(new byte[0], Errors.UNKNOWN_MEMBER_ID); + } else if (generationId != group.generationId()) { + responseCallback.accept(new byte[0], Errors.ILLEGAL_GENERATION); + } else { + switch (group.currentState()) { + case Empty: + case Dead: + responseCallback.accept(new byte[0], Errors.UNKNOWN_MEMBER_ID); + break; + + case PreparingRebalance: + responseCallback.accept(new byte[0], Errors.REBALANCE_IN_PROGRESS); + break; + + case CompletingRebalance: + group.get(memberId).awaitingSyncCallback(responseCallback); + + // if this is the leader, then we can attempt to persist state and transition to stable + if (group.isLeader(memberId)) { + log.info("Assignment received from leader for group {} for generation {}", + group.groupId(), group.generationId()); + + // fill any missing members with an empty assignment + Set missing = Sets.difference(group.allMembers(), groupAssignment.keySet()); + Map assignment = new HashMap<>(); + assignment.putAll(groupAssignment); + assignment.putAll( + missing.stream() + .collect(Collectors.toMap( + k -> k, + k -> new byte[0] + )) + ); + + groupManager.storeGroup(group, assignment).thenApply(error -> { + return group.inLock(() -> { + // another member may have joined the group while we were awaiting this callback, + // so we must ensure we are still in the CompletingRebalance state and the same + // generation when it gets invoked. if we have transitioned to another state, + // then do nothing + if (group.is(CompletingRebalance) && generationId == group.generationId()) { + if (error != Errors.NONE) { + resetAndPropagateAssignmentError(group, error); + maybePrepareRebalance(group); + } else { + setAndPropagateAssignment(group, assignment); + group.transitionTo(Stable); + } + } + return null; + }); + }); + } + break; + + case Stable: + // if the group is stable, we just return the current assignment + MemberMetadata memberMetadata = group.get(memberId); + responseCallback.accept(memberMetadata.assignment(), Errors.NONE); + completeAndScheduleNextHeartbeatExpiration(group, group.get(memberId)); + break; + + default: + throw new IllegalStateException("Should not reach here"); + } + } + return null; + }); + } + + public CompletableFuture handleLeaveGroup( + String groupId, + String memberId + ) { + return validateGroupStatus(groupId, ApiKeys.LEAVE_GROUP).map(error -> + CompletableFuture.completedFuture(error) + ).orElseGet(() -> { + return groupManager.getGroup(groupId).map(group -> { + return group.inLock(() -> { + if (group.is(Dead) || !group.has(memberId)) { + return CompletableFuture.completedFuture(Errors.UNKNOWN_MEMBER_ID); + } else { + MemberMetadata member = group.get(memberId); + removeHeartbeatForLeavingMember(group, member); + if (log.isDebugEnabled()) { + log.debug("Member {} in group {} has left, removing it from the group", + member.memberId(), group.groupId()); + } + removeMemberAndUpdateGroup(group, member); + return CompletableFuture.completedFuture(Errors.NONE); + } + }); + }).orElseGet(() -> { + // if the group is marked as dead, it means some other thread has just removed the group + // from the coordinator metadata; this is likely that the group has migrated to some other + // coordinator OR the group is in a transient unstable phase. Let the consumer to retry + // joining without specified consumer id, + return CompletableFuture.completedFuture(Errors.UNKNOWN_MEMBER_ID); + }); + }); + } + + public Map handleDeleteGroups(Set groupIds) { + Map groupErrors = new HashMap<>(); + List groupsEligibleForDeletion = new ArrayList<>(); + + groupIds.forEach(groupId -> { + validateGroupStatus(groupId, ApiKeys.DELETE_GROUPS).map(error -> + groupErrors.put(groupId, error) + ).orElseGet(() -> groupManager.getGroup(groupId).map(group -> { + return group.inLock(() -> { + switch(group.currentState()) { + case Dead: + if (groupManager.groupNotExists(groupId)) { + groupErrors.put(groupId, Errors.GROUP_ID_NOT_FOUND); + } else { + groupErrors.put(groupId, Errors.NOT_COORDINATOR); + } + break; + case Empty: + group.transitionTo(Dead); + groupsEligibleForDeletion.add(group); + break; + default: + groupErrors.put(groupId, Errors.NON_EMPTY_GROUP); + break; + } + return Errors.NONE; + }); + }).orElseGet(() -> { + Errors error; + if (groupManager.groupNotExists(groupId)) { + error = Errors.GROUP_ID_NOT_FOUND; + } else { + error = Errors.NOT_COORDINATOR; + } + groupErrors.put(groupId, error); + return Errors.NONE; + })); + }); + + if (!groupsEligibleForDeletion.isEmpty()) { + // TODO: + /// val offsetsRemoved = groupManager.cleanupGroupMetadata(groupsEligibleForDeletion, _.removeAllOffsets()) + /// groupErrors ++= groupsEligibleForDeletion.map(_.groupId -> Errors.NONE).toMap + /// info(s"The following groups were deleted: ${groupsEligibleForDeletion.map(_.groupId).mkString(", ")}. " + + /// s"A total of $offsetsRemoved offsets were removed.") + } + + return groupErrors; + } + + public CompletableFuture handleHeartbeat(String groupId, + String memberId, + int generationId) { + return validateGroupStatus(groupId, ApiKeys.HEARTBEAT).map(error -> { + if (error == Errors.COORDINATOR_LOAD_IN_PROGRESS) { + // the group is still loading, so respond just blindly + return CompletableFuture.completedFuture(Errors.NONE); + } else { + return CompletableFuture.completedFuture(error); + } + }).orElseGet(() -> groupManager.getGroup(groupId).map(group -> + group.inLock(() -> { + switch(group.currentState()) { + case Dead: + // if the group is marked as dead, it means some other thread has just removed the group + // from the coordinator metadata; this is likely that the group has migrated to some other + // coordinator OR the group is in a transient unstable phase. Let the member retry + // joining without the specified member id, + return CompletableFuture.completedFuture(Errors.UNKNOWN_MEMBER_ID); + + case Empty: + return CompletableFuture.completedFuture(Errors.UNKNOWN_MEMBER_ID); + + case CompletingRebalance: + if (!group.has(memberId)) { + return CompletableFuture.completedFuture(Errors.UNKNOWN_MEMBER_ID); + } else { + return CompletableFuture.completedFuture(Errors.REBALANCE_IN_PROGRESS); + } + + case PreparingRebalance: + if (!group.has(memberId)) { + return CompletableFuture.completedFuture(Errors.UNKNOWN_MEMBER_ID); + } else if (generationId != group.generationId()) { + return CompletableFuture.completedFuture(Errors.ILLEGAL_GENERATION); + } else { + MemberMetadata member = group.get(memberId); + completeAndScheduleNextHeartbeatExpiration(group, member); + return CompletableFuture.completedFuture(Errors.REBALANCE_IN_PROGRESS); + } + + case Stable: + if (!group.has(memberId)) { + return CompletableFuture.completedFuture(Errors.UNKNOWN_MEMBER_ID); + } else if (generationId != group.generationId()) { + return CompletableFuture.completedFuture(Errors.ILLEGAL_GENERATION); + } else { + MemberMetadata member = group.get(memberId); + completeAndScheduleNextHeartbeatExpiration(group, member); + return CompletableFuture.completedFuture(Errors.NONE); + } + + default: + return CompletableFuture.completedFuture(Errors.NONE); + } + }) + ).orElseGet(() -> + CompletableFuture.completedFuture(Errors.UNKNOWN_MEMBER_ID) + )); + } + + KeyValue> handleListGroups() { + if (!isActive.get()) { + return new KeyValue<>(Errors.COORDINATOR_NOT_AVAILABLE, new ArrayList<>()); + } else { + Errors errors; + if (groupManager.isLoading()) { + errors = Errors.COORDINATOR_LOAD_IN_PROGRESS; + } else { + errors = Errors.NONE; + } + List overviews = new ArrayList<>(); + groupManager.currentGroups().forEach(group -> overviews.add(group.overview())); + return new KeyValue<>( + errors, + overviews + ); + } + } + + KeyValue handleDescribeGroup(String groupId) { + return validateGroupStatus(groupId, ApiKeys.DESCRIBE_GROUPS).map(error -> + new KeyValue<>(error, GroupCoordinator.EmptyGroup) + ).orElseGet(() -> + groupManager.getGroup(groupId) + .map(group -> + group.inLock(() -> new KeyValue(Errors.NONE, group.summary()) + )) + .orElseGet(() -> new KeyValue<>(Errors.NONE, GroupCoordinator.DeadGroup)) + ); + } + + private Optional validateGroupStatus(String groupId, + ApiKeys api) { + if (isValidGroupId(groupId, api)) { + return Optional.of(Errors.INVALID_GROUP_ID); + } else if (isActive.get()) { + return Optional.of(Errors.COORDINATOR_NOT_AVAILABLE); + } else if (groupManager.isGroupLoading(groupId)) { + return Optional.of(Errors.COORDINATOR_LOAD_IN_PROGRESS); + } else if (groupManager.isGroupLocal(groupId)) { + return Optional.of(Errors.NOT_COORDINATOR); + } else { + return Optional.empty(); + } + } + + private void setAndPropagateAssignment(GroupMetadata group, + Map assignment) { + checkState(group.is(CompletingRebalance)); + group.allMemberMetadata().forEach(member -> member.assignment(assignment.get(member.memberId()))); + propagateAssignment(group, Errors.NONE); + } + + private void resetAndPropagateAssignmentError(GroupMetadata group, + Errors error) { + checkState(group.is(CompletingRebalance)); + group.allMemberMetadata().forEach(m -> m.assignment(new byte[0])); + propagateAssignment(group, error); + } + + private void propagateAssignment(GroupMetadata group, Errors error) { + for (MemberMetadata member : group.allMemberMetadata()) { + if (member.awaitingSyncCallback() != null) { + member.awaitingSyncCallback().accept(member.assignment(), error); + member.awaitingSyncCallback(null); + + // reset the session timeout for members after propagating the member's assignment. + // This is because if any member's session expired while we were still awaiting either + // the leader sync group or the storage callback, its expiration will be ignored and no + // future heartbeat expectations will not be scheduled. + completeAndScheduleNextHeartbeatExpiration(group, member) + } + } + } + + private JoinGroupResult joinError(String memberId, Errors error) { + return new JoinGroupResult( + Collections.emptyMap(), + memberId, + 0, + GroupCoordinator.NoProtocol, + GroupCoordinator.NoLeader, + error); + } + + /** + * Complete existing DelayedHeartbeats for the given member and schedule the next one + */ + private void completeAndScheduleNextHeartbeatExpiration(GroupMetadata group, MemberMetadata member) { + // complete current heartbeat expectation + member.latestHeartbeat(time.milliseconds()); + val memberKey = MemberKey(member.groupId, member.memberId) + heartbeatPurgatory.checkAndComplete(memberKey) + + // reschedule the next heartbeat expiration deadline + long newHeartbeatDeadline = member.latestHeartbeat() + member.sessionTimeoutMs(); + val delayedHeartbeat = new DelayedHeartbeat(this, group, member, newHeartbeatDeadline, member.sessionTimeoutMs()); + heartbeatPurgatory.tryCompleteElseWatch(delayedHeartbeat, Seq(memberKey)) + } + + private void removeHeartbeatForLeavingMember(GroupMetadata group, + MemberMetadata member) { + member.isLeaving(true); + val memberKey = MemberKey(member.groupId, member.memberId) + heartbeatPurgatory.checkAndComplete(memberKey); + } + + private CompletableFuture addMemberAndRebalance( + int rebalanceTimeoutMs, + int sessionTimeoutMs, + String clientId, + String clientHost, + String protocolType, + Map protocols, + GroupMetadata group + ) { + String memberId = clientId + "-" + group.generateMemberIdSuffix(); + MemberMetadata member = new MemberMetadata( + memberId, + group.groupId(), + clientId, + clientHost, + rebalanceTimeoutMs, + sessionTimeoutMs, + protocolType, + protocols); + CompletableFuture joinFuture = new CompletableFuture<>(); + member.awaitingJoinCallback(joinFuture); + // update the newMemberAdded flag to indicate that the join group can be further delayed + if (group.is(PreparingRebalance) && group.generationId() == 0) { + group.newMemberAdded(true); + } + + group.add(member); + maybePrepareRebalance(group); + return joinFuture; + } + + private CompletableFuture updateMemberAndRebalance( + GroupMetadata group, + MemberMetadata member, + Map protocols + ) { + CompletableFuture resultFuture = new CompletableFuture<>(); + member.supportedProtocols(protocols); + member.awaitingJoinCallback(resultFuture); + maybePrepareRebalance(group); + return resultFuture; + } + + private void maybePrepareRebalance(GroupMetadata group) { + group.inLock(() -> { + if (group.canReblance()) { + prepareRebalance(group); + } + return null; + }); + } + + + + private void prepareRebalance(GroupMetadata group) { + // if any members are awaiting sync, cancel their request and have them rejoin + if (group.is(CompletingRebalance)) + resetAndPropagateAssignmentError(group, Errors.REBALANCE_IN_PROGRESS) + + val delayedRebalance = if (group.is(Empty)) + new InitialDelayedJoin(this, + joinPurgatory, + group, + groupConfig.groupInitialRebalanceDelayMs, + groupConfig.groupInitialRebalanceDelayMs, + max(group.rebalanceTimeoutMs - groupConfig.groupInitialRebalanceDelayMs, 0)) + else + new DelayedJoin(this, group, group.rebalanceTimeoutMs) + + group.transitionTo(PreparingRebalance) + + info(s"Preparing to rebalance group ${group.groupId} with old generation ${group.generationId} " + + s"(${Topic.GROUP_METADATA_TOPIC_NAME}-${partitionFor(group.groupId)})") + + val groupKey = GroupKey(group.groupId) + joinPurgatory.tryCompleteElseWatch(delayedRebalance, Seq(groupKey)) + } + + private void removeMemberAndUpdateGroup(GroupMetadata group, + MemberMetadata member) { + group.remove(member.memberId()); + switch (group.currentState()) { + case Dead: + case Empty: + break; + case Stable: + case CompletingRebalance: + maybePrepareRebalance(group); + break; + case PreparingRebalance: + // joinPurgatory.checkAndComplete(GroupKey(group.groupId)) + break; + default: + break; + } + } + + boolean tryCompleteJoin(GroupMetadata group, + Supplier forceComplete) { + return group.inLock(() -> { + if (group.notYetRejoinedMembers().isEmpty()) { + return forceComplete.get(); + } else { + return false; + } + }); + } + + void onExpireJoin() { + // TODO: add metrics for restabilize timeouts + } + + void onCompleteJoin(GroupMetadata group) { + group.inLock(() -> { + // remove any members who haven't joined the group yet + group.notYetRejoinedMembers().forEach(failedMember -> { + removeHeartbeatForLeavingMember(group, failedMember) + group.remove(failedMember.memberId()); + // TODO: cut the socket connection to the client + }); + + if (!group.is(Dead)) { + group.initNextGeneration(); + if (group.is(Empty)) { + log.info("Group {} with generation {} is now empty {}-{}", + group.groupId(), group.generationId(), + Topic.GROUP_METADATA_TOPIC_NAME, groupManager.partitionFor(group.groupId())); + + groupManager.storeGroup(group, Collections.emptyMap(), error => { + if (error != Errors.NONE) { + // we failed to write the empty group metadata. If the broker fails before another rebalance, + // the previous generation written to the log will become active again (and most likely timeout). + // This should be safe since there are no active members in an empty generation, so we just warn. + warn(s"Failed to write empty metadata for group ${group.groupId}: ${error.message}") + } + }) + } else { + info(s"Stabilized group ${group.groupId} generation ${group.generationId} " + + s"(${Topic.GROUP_METADATA_TOPIC_NAME}-${partitionFor(group.groupId)})") + + // trigger the awaiting join group response callback for all the members after rebalancing + for (MemberMetadata member : group.allMemberMetadata()) { + Objects.requireNonNull(member.awaitingJoinCallback()); + JoinGroupRequest joinResult; + if (group.isLeader(member.memberId())) { + joinResult = new JoinGroupResult( + group.currentMemberMetadata + ); + } else { + joinResult = new JoinGroupResult( + Collections.emptyMap() + ); + } + = JoinGroupResult( + members = if (group.isLeader(member.memberId)) { + group.currentMemberMetadata + } else { + Map.empty + }, + memberId = member.memberId, + generationId = group.generationId, + subProtocol = group.protocolOrNull, + leaderId = group.leaderOrNull, + error = Errors.NONE) + + member.awaitingJoinCallback(joinResult) + member.awaitingJoinCallback = null + completeAndScheduleNextHeartbeatExpiration(group, member) + } + } + } + }); + } + + boolean tryCompleteHeartbeat(GroupMetadata group, + MemberMetadata member, + long heartbeatDeadline, + Supplier forceComplete) { + return group.inLock(() -> { + if (shouldKeepMemberAlive(member, heartbeatDeadline) + || member.isLeaving()) { + return forceComplete.get(); + } else { + return false; + } + }); + } + + void onExpireHeartbeat(GroupMetadata group, + MemberMetadata member, + long heartbeatDeadline) { + group.inLock(() -> { + if (!shouldKeepMemberAlive(member, heartbeatDeadline)) { + log.info("Member {} in group {} has failed, removing it from the group", + member.memberId(), group.groupId()); + removeMemberAndUpdateGroup(group, member); + } + return null; + }); + } + + void onCompleteHeartbeat() { + // TODO: add metrics for complete heartbeats + } + + private boolean shouldKeepMemberAlive(MemberMetadata member, + long heartbeatDeadline) { + return member.awaitingJoinCallback() != null || + member.awaitingSyncCallback() != null || + member.latestHeartbeat() + member.sessionTimeoutMs() > heartbeatDeadline; + } } diff --git a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadata.java b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadata.java index 021c5ea274..7aad590bf3 100644 --- a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadata.java +++ b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadata.java @@ -15,6 +15,8 @@ import static com.google.common.base.Preconditions.checkArgument; import static com.google.common.base.Preconditions.checkState; +import static io.streamnative.kop.coordinator.group.GroupState.Dead; +import static io.streamnative.kop.coordinator.group.GroupState.PreparingRebalance; import com.google.common.base.MoreObjects; import com.google.common.base.MoreObjects.ToStringHelper; @@ -32,6 +34,7 @@ import java.util.Objects; import java.util.Optional; import java.util.Set; +import java.util.UUID; import java.util.concurrent.locks.ReentrantLock; import java.util.stream.Collectors; import javax.annotation.concurrent.NotThreadSafe; @@ -62,20 +65,20 @@ class GroupMetadata { static { validPreviousStates.put( - GroupState.Dead, + Dead, Sets.newHashSet( GroupState.Stable, - GroupState.PreparingRebalance, + PreparingRebalance, GroupState.CompletingRebalance, GroupState.Empty, - GroupState.Dead + Dead ) ); validPreviousStates.put( GroupState.CompletingRebalance, Sets.newHashSet( - GroupState.PreparingRebalance + PreparingRebalance ) ); @@ -87,7 +90,7 @@ class GroupMetadata { ); validPreviousStates.put( - GroupState.PreparingRebalance, + PreparingRebalance, Sets.newHashSet( GroupState.Stable, GroupState.CompletingRebalance, @@ -98,7 +101,7 @@ class GroupMetadata { validPreviousStates.put( GroupState.Empty, Sets.newHashSet( - GroupState.PreparingRebalance + PreparingRebalance ) ); } @@ -151,6 +154,7 @@ static class GroupSummary { private long generationId = 0L; private Optional leaderId = Optional.empty(); private Optional protocol = Optional.empty(); + private boolean newMemberAdded = false; // state management private final Map members = new HashMap<>(); @@ -160,6 +164,14 @@ static class GroupSummary { this.state = initialState; } + public String generateMemberIdSuffix() { + return UUID.randomUUID().toString(); + } + + public void newMemberAdded(boolean newMemberAdded) { + this.newMemberAdded = newMemberAdded; + } + public T inLock(Supplier supplier) { return CoreUtils.inLock(lock, supplier); } @@ -180,6 +192,10 @@ public long generationId() { return generationId; } + public Set allMembers() { + return members.keySet(); + } + public List allMemberMetadata() { return members.values().stream().collect(Collectors.toList()); } @@ -196,6 +212,10 @@ public boolean has(String memberId) { return members.containsKey(memberId); } + public MemberMetadata get(String memberId) { + return members.get(memberId); + } + public boolean isLeader(String memberId) { return Objects.equals(leaderId.orElse(null), memberId); } @@ -272,7 +292,7 @@ public void remove(String memberId) { } public boolean canReblance() { - return validPreviousStates.get(GroupState.PreparingRebalance).contains(state); + return validPreviousStates.get(PreparingRebalance).contains(state); } public void transitionTo(GroupState groupState) { @@ -309,6 +329,17 @@ public String selectProtocol() { .orElse(null); } + public Map currentMemberMetadata() { + if (is(Dead) || is(PreparingRebalance)) { + throw new IllegalStateException("Cannot obtain member metadata for group in state " + state); + } + return members.entrySet().stream() + .collect(Collectors.toMap( + e -> e.getKey(), + e -> e.getValue().metadata(protocol.get()) + )); + } + public GroupSummary summary() { if (is(GroupState.Stable)) { String protocol = protocolOrNull(); diff --git a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataConstants.java b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataConstants.java index fa6b06ebff..04b18ab34b 100644 --- a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataConstants.java +++ b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataConstants.java @@ -8,10 +8,14 @@ import static org.apache.kafka.common.protocol.types.Type.STRING; import com.google.common.collect.Lists; +import io.streamnative.kop.coordinator.group.GroupMetadataManager.BaseKey; +import io.streamnative.kop.coordinator.group.GroupMetadataManager.GroupMetadataKey; import java.nio.ByteBuffer; +import java.util.HashMap; import java.util.List; import java.util.Map; import java.util.stream.Collectors; +import java.util.stream.Stream; import lombok.val; import org.apache.kafka.common.KafkaException; import org.apache.kafka.common.protocol.types.ArrayOf; @@ -19,6 +23,7 @@ import org.apache.kafka.common.protocol.types.Field; import org.apache.kafka.common.protocol.types.Schema; import org.apache.kafka.common.protocol.types.Struct; +import org.apache.kafka.common.utils.Utils; import org.apache.pulsar.common.schema.KeyValue; /** @@ -260,6 +265,95 @@ static byte[] groupMetadataValue(GroupMetadata groupMetadata, return byteBuffer.array(); } + /** + * Decodes the offset messages' key + */ + static BaseKey readMessageKey(ByteBuffer buffer) { + short version = buffer.getShort(); + Schema keySchema = schemaForKey(version); + Struct key = keySchema.read(buffer); + + if (version <= CURRENT_OFFSET_KEY_SCHEMA_VERSION) { + // version 0 and 1 refer to offset + throw new UnsupportedOperationException(); + } else if (version == CURRENT_GROUP_KEY_SCHEMA_VERSION) { + // version 2 refers to group + String group = key.getString(GROUP_KEY_GROUP_FIELD); + + return new GroupMetadataKey(version, group); + } else { + throw new IllegalStateException("Unknown version " + version + " for group metadata message"); + } + } + + static GroupMetadata readGroupMessageValue(String groupId, + ByteBuffer buffer) { + if (null == buffer) { // tombstone + return null; + } + + short version = buffer.getShort(); + Schema valueSchema = schemaForGroup(version); + Struct value = valueSchema.read(buffer); + + if (version == 0 || version == 1) { + int generationId = value.getInt(GENERATION_KEY); + String protocolType = value.getString(PROTOCOL_TYPE_KEY); + String protocol = value.getString(PROTOCOL_KEY); + String leaderId = value.getString(LEADER_KEY); + Object[] memberMetadataArray = value.getArray(MEMBERS_KEY); + GroupState initialState; + if (memberMetadataArray.length == 0) { + initialState = GroupState.Empty; + } else { + initialState = GroupState.Stable; + } + + List members = Lists.newArrayList(memberMetadataArray) + .stream() + .map(memberMetadataObj -> { + Struct memberMetadata = (Struct) memberMetadataObj; + String memberId = memberMetadata.getString(MEMBER_ID_KEY); + String clientId = memberMetadata.getString(CLIENT_ID_KEY); + String clientHost = memberMetadata.getString(CLIENT_HOST_KEY); + int sessionTimeout = memberMetadata.getInt(SESSION_TIMEOUT_KEY); + int rebalanceTimeout; + if (version == 0) { + rebalanceTimeout = sessionTimeout; + } else { + rebalanceTimeout = memberMetadata.getInt(REBALANCE_TIMEOUT_KEY); + } + ByteBuffer subscription = memberMetadata.getBytes(SUBSCRIPTION_KEY); + byte[] subscriptionData = new byte[subscription.remaining()]; + subscription.get(subscriptionData); + Map protocols = new HashMap<>(); + protocols.put(protocol, subscriptionData); + return new MemberMetadata( + memberId, + groupId, + clientId, + clientHost, + rebalanceTimeout, + sessionTimeout, + protocolType, + protocols + ); + }).collect(Collectors.toList()); + + return GroupMetadata.loadGroup( + groupId, + initialState, + generationId, + protocolType, + protocol, + leaderId, + members + ); + } else { + throw new IllegalStateException("Unknown group metadata message version"); + } + } + private GroupMetadataConstants() {} } diff --git a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataManager.java b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataManager.java index 5d72e725dd..22ff1cf731 100644 --- a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataManager.java +++ b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataManager.java @@ -1,12 +1,16 @@ package io.streamnative.kop.coordinator.group; -import static io.streamnative.kop.coordinator.group.GroupMetadataConstants.CURRENT_GROUP_KEY_SCHEMA; -import static io.streamnative.kop.coordinator.group.GroupMetadataConstants.CURRENT_GROUP_KEY_SCHEMA_VERSION; -import static io.streamnative.kop.coordinator.group.GroupMetadataConstants.GROUP_KEY_GROUP_FIELD; +import static io.streamnative.kop.coordinator.group.GroupMetadataConstants.CURRENT_GROUP_VALUE_SCHEMA_VERSION; +import static io.streamnative.kop.coordinator.group.GroupMetadataConstants.groupMetadataKey; +import static io.streamnative.kop.coordinator.group.GroupMetadataConstants.groupMetadataValue; +import static io.streamnative.kop.coordinator.group.GroupMetadataConstants.readGroupMessageValue; +import static io.streamnative.kop.coordinator.group.GroupMetadataConstants.readMessageKey; import static io.streamnative.kop.utils.CoreUtils.inLock; import static java.nio.charset.StandardCharsets.UTF_8; +import static org.apache.kafka.common.internals.Topic.GROUP_METADATA_TOPIC_NAME; import java.nio.ByteBuffer; +import java.util.HashMap; import java.util.HashSet; import java.util.Map; import java.util.Optional; @@ -19,18 +23,22 @@ import java.util.function.Consumer; import lombok.Data; import lombok.experimental.Accessors; +import lombok.extern.slf4j.Slf4j; import org.apache.bookkeeper.common.hash.Murmur3; import org.apache.bookkeeper.common.util.MathUtils; -import org.apache.bookkeeper.common.util.OrderedScheduler; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.protocol.Errors; -import org.apache.kafka.common.protocol.types.Struct; import org.apache.kafka.common.record.CompressionType; import org.apache.kafka.common.requests.ApiVersionsResponse.ApiVersion; +import org.apache.kafka.common.utils.Time; +import org.apache.pulsar.client.api.MessageId; +import org.apache.pulsar.client.api.Producer; +import org.apache.pulsar.client.api.Reader; /** * Manager to manage a coordination group. */ +@Slf4j class GroupMetadataManager { /** @@ -46,7 +54,7 @@ interface BaseKey { */ @Data @Accessors(fluent = true) - static class GroupMetadataKey { + static class GroupMetadataKey implements BaseKey { private final short version; private final String key; @@ -101,15 +109,19 @@ public String toString() { private final Set ownedPartitions = new HashSet<>(); /* shutting down flag */ private final AtomicBoolean shuttingDown = new AtomicBoolean(false); - /* single-thread scheduler to handle offset/group metadata cache loading and unloading */ - private OrderedScheduler scheduler; private final String logIdent; private final int groupMetadataTopicPartitionCount; + private final Producer metadataTopicProducer; + private final Reader metadataTopicReader; + private final Time time; GroupMetadataManager(int brokerId, ApiVersion interBrokerProtocolVersion, OffsetConfig config, - int groupMetadataTopicPartitionCount) { + int groupMetadataTopicPartitionCount, + Producer metadataTopicProducer, + Reader metadataTopicConsumer, + Time time) { this.brokerId = brokerId; this.interBrokerProtocolVersion = interBrokerProtocolVersion; this.config = config; @@ -117,13 +129,9 @@ public String toString() { this.groupMetadataCache = new ConcurrentHashMap<>(); this.logIdent = String.format("[GroupMetadataManager brokerId=%d]", brokerId); this.groupMetadataTopicPartitionCount = groupMetadataTopicPartitionCount; - } - - public void startup() { - this.scheduler = OrderedScheduler.newSchedulerBuilder() - .name("group-metadata-manager") - .numThreads(1) - .build(); + this.metadataTopicProducer = metadataTopicProducer; + this.metadataTopicReader = metadataTopicConsumer; + this.time = time; } public Iterable currentGroups() { @@ -197,7 +205,198 @@ public GroupMetadata addGroup(GroupMetadata group) { public CompletableFuture storeGroup(GroupMetadata group, Map groupAssignment) { + long timestamp = time.milliseconds(); + byte[] key = groupMetadataKey(group.groupId()); + byte[] value = groupMetadataValue( + group, groupAssignment, CURRENT_GROUP_VALUE_SCHEMA_VERSION); + + return metadataTopicProducer.newMessage() + .keyBytes(key) + .value(value) + .eventTime(timestamp) + .sendAsync() + .thenApply(msgId -> Errors.NONE) + .exceptionally(cause -> Errors.COORDINATOR_NOT_AVAILABLE); + } + + public CompletableFuture scheduleLoadGroupAndOffsets(int offsetsPartition, + Consumer onGroupLoaded) { + TopicPartition topicPartition = new TopicPartition( + GROUP_METADATA_TOPIC_NAME, offsetsPartition + ); + if (addLoadingPartition(offsetsPartition)) { + log.info("Scheduling loading of offsets and group metadata from {}", topicPartition); + long startMs = time.milliseconds(); + return metadataTopicProducer.newMessage() + .value(new byte[0]) + .eventTime(time.milliseconds()) + .sendAsync() + .thenCompose(lastMessageId -> + doLoadGroupsAndOffsets(metadataTopicReader, lastMessageId, onGroupLoaded)) + .whenComplete((ignored, cause) -> { + if (null == cause) { + log.info("Finished loading offsets and group metadata from {} in {} milliseconds", + topicPartition, time.milliseconds() - startMs); + } else { + log.error("Error loading offsets from {}", topicPartition, cause); + } + inLock(partitionLock, () -> { + ownedPartitions.add(topicPartition.partition()); + loadingPartitions.remove(topicPartition.partition()); + return null; + }); + }); + } else { + log.info("Already loading offsets and group metadata from {}", topicPartition); + return CompletableFuture.completedFuture(null); + } + } + + private CompletableFuture doLoadGroupsAndOffsets( + Reader metadataConsumer, + MessageId endMessageId, + Consumer onGroupLoaded + ) { + final Map loadedGroups = new HashMap<>(); + final Set removedGroups = new HashSet<>(); + final CompletableFuture resultFuture = new CompletableFuture<>(); + + loadNextMetadataMessage( + metadataConsumer, + endMessageId, + resultFuture, + onGroupLoaded, + loadedGroups, + removedGroups); + + return resultFuture; + } + + private void loadNextMetadataMessage(Reader metadataConsumer, + MessageId endMessageId, + CompletableFuture resultFuture, + Consumer onGroupLoaded, + Map loadedGroups, + Set removedGroups) { + metadataConsumer.readNextAsync().whenComplete((message, cause) -> { + if (null != cause) { + resultFuture.completeExceptionally(cause); + return; + } + + if (message.getMessageId().compareTo(endMessageId) >= 0) { + // reach the end of partition + processLoadedAndRemovedGroups( + resultFuture, + onGroupLoaded, + loadedGroups, + removedGroups + ); + return; + } + if (!message.hasKey()) { + // the messages without key are placeholders + loadNextMetadataMessage( + metadataConsumer, + endMessageId, + resultFuture, + onGroupLoaded, + loadedGroups, + removedGroups + ); + return; + } + + BaseKey baseKey = readMessageKey(ByteBuffer.wrap(message.getKeyBytes())); + if (baseKey instanceof GroupMetadataKey) { + // load group metadata + GroupMetadataKey gmKey = (GroupMetadataKey) baseKey; + String groupId = gmKey.key(); + byte[] data = message.getValue(); + if (null == data || data.length == 0) { + // null value + loadedGroups.remove(groupId); + removedGroups.add(groupId); + } else { + GroupMetadata groupMetadata = readGroupMessageValue( + groupId, + ByteBuffer.wrap(message.getValue()) + ); + if (null != groupMetadata) { + removedGroups.remove(groupId); + loadedGroups.put(groupId, groupMetadata); + } else { + loadedGroups.remove(groupId); + removedGroups.add(groupId); + } + } + + loadNextMetadataMessage( + metadataConsumer, + endMessageId, + resultFuture, + onGroupLoaded, + loadedGroups, + removedGroups + ); + } else { + resultFuture.completeExceptionally( + new IllegalStateException("Unexpected message key " + + baseKey + " while loading offsets and group metadata")); + } + + }); + } + + private void processLoadedAndRemovedGroups(CompletableFuture resultFuture, + Consumer onGroupLoaded, + Map loadedGroups, + Set removedGroups) { + try { + loadedGroups.values().forEach(group -> { + loadGroup(group); + onGroupLoaded.accept(group); + }); + + removedGroups.forEach(groupId -> { + // TODO: add offsets later + }); + resultFuture.complete(null); + } catch (RuntimeException re) { + resultFuture.completeExceptionally(re); + } + } + + private void loadGroup(GroupMetadata group) { + GroupMetadata currentGroup = addGroup(group); + if (group != currentGroup) { + log.debug("Attempt to load group {} from log with generation {} failed " + + "because there is already a cached group with generation {}", + group.groupId(), group.generationId(), currentGroup.generationId()); + } + } + + /** + * Add the partition into the owned list + * + * NOTE: this is for test only + */ + private void addPartitionOwnership(int partition) { + inLock(partitionLock, () -> { + ownedPartitions.add(partition); + return null; + }); + } + + /** + * Add a partition to the loading partitions set. Return true if the partition was not + * already loading. + * + * Visible for testing + */ + boolean addLoadingPartition(int partition) { + return inLock(partitionLock, () -> loadingPartitions.add(partition)); } diff --git a/src/main/java/io/streamnative/kop/coordinator/group/MemberMetadata.java b/src/main/java/io/streamnative/kop/coordinator/group/MemberMetadata.java index bb6f728768..7c27ff3604 100644 --- a/src/main/java/io/streamnative/kop/coordinator/group/MemberMetadata.java +++ b/src/main/java/io/streamnative/kop/coordinator/group/MemberMetadata.java @@ -22,8 +22,8 @@ import java.util.Map; import java.util.Optional; import java.util.Set; +import java.util.concurrent.CompletableFuture; import java.util.function.BiConsumer; -import java.util.function.Consumer; import javax.annotation.concurrent.NotThreadSafe; import lombok.Data; import lombok.Getter; @@ -89,14 +89,18 @@ public static class MemberSummary { private final Map supportedProtocols; private byte[] assignment = new byte[0]; - private Consumer awaitingJoinCallback = null; + private CompletableFuture awaitingJoinCallback = null; private BiConsumer awaitingSyncCallback = null; private long latestHeartbeat = -1L; private boolean isLeaving = false; public Set protocols() { return supportedProtocols.keySet(); + } + public void supportedProtocols(Map supportedProtocols) { + this.supportedProtocols.clear(); + this.supportedProtocols.putAll(supportedProtocols); } public byte[] metadata(String protocol) { diff --git a/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataManagerTest.java b/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataManagerTest.java new file mode 100644 index 0000000000..296c8ebb2f --- /dev/null +++ b/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataManagerTest.java @@ -0,0 +1,10 @@ +functionspackage io.streamnative.kop.coordinator.group; + +/** + * Unit test {@link GroupMetadataManager}. + */ +public class GroupMetadataManagerTest { + + + +} diff --git a/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataTest.java b/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataTest.java index 7b73e62f2a..adb727f3df 100644 --- a/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataTest.java +++ b/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataTest.java @@ -28,6 +28,7 @@ import java.util.LinkedHashMap; import java.util.Map; import java.util.Set; +import java.util.concurrent.CompletableFuture; import lombok.val; import org.junit.Before; import org.junit.Test; @@ -351,7 +352,7 @@ public void testInitNextGeneration() { protocols); group.transitionTo(PreparingRebalance); - member.awaitingJoinCallback(e -> {}); + member.awaitingJoinCallback(new CompletableFuture<>()); group.add(member); assertEquals(0, group.generationId()); From e914e3fa37221e9fea4c9b6bb38a199c4196e789 Mon Sep 17 00:00:00 2001 From: Sijie Guo Date: Sat, 3 Aug 2019 16:56:53 +0800 Subject: [PATCH 4/8] Port Kafka timer classes *Motivation* Need to port Kafka's coordinator algorithm. It requires using Kafka's timer and deplayed operations. --- .../kop/utils/timer/SystemTimer.java | 181 +++++++++++++++ .../streamnative/kop/utils/timer/Timer.java | 51 +++++ .../kop/utils/timer/TimerTask.java | 50 +++++ .../kop/utils/timer/TimerTaskList.java | 209 ++++++++++++++++++ .../kop/utils/timer/TimingWheel.java | 206 +++++++++++++++++ .../kop/utils/timer/package-info.java | 19 ++ .../io/streamnative/kop/utils/MockTime.java | 111 ++++++++++ .../kop/utils/timer/MockTimer.java | 85 +++++++ .../kop/utils/timer/TimerTaskListTest.java | 108 +++++++++ .../kop/utils/timer/TimerTest.java | 161 ++++++++++++++ 10 files changed, 1181 insertions(+) create mode 100644 src/main/java/io/streamnative/kop/utils/timer/SystemTimer.java create mode 100644 src/main/java/io/streamnative/kop/utils/timer/Timer.java create mode 100644 src/main/java/io/streamnative/kop/utils/timer/TimerTask.java create mode 100644 src/main/java/io/streamnative/kop/utils/timer/TimerTaskList.java create mode 100644 src/main/java/io/streamnative/kop/utils/timer/TimingWheel.java create mode 100644 src/main/java/io/streamnative/kop/utils/timer/package-info.java create mode 100644 src/test/java/io/streamnative/kop/utils/MockTime.java create mode 100644 src/test/java/io/streamnative/kop/utils/timer/MockTimer.java create mode 100644 src/test/java/io/streamnative/kop/utils/timer/TimerTaskListTest.java create mode 100644 src/test/java/io/streamnative/kop/utils/timer/TimerTest.java diff --git a/src/main/java/io/streamnative/kop/utils/timer/SystemTimer.java b/src/main/java/io/streamnative/kop/utils/timer/SystemTimer.java new file mode 100644 index 0000000000..9056db31fc --- /dev/null +++ b/src/main/java/io/streamnative/kop/utils/timer/SystemTimer.java @@ -0,0 +1,181 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.utils.timer; + +import com.google.common.util.concurrent.ThreadFactoryBuilder; +import io.streamnative.kop.utils.timer.TimerTaskList.TimerTaskEntry; +import java.util.Objects; +import java.util.concurrent.DelayQueue; +import java.util.concurrent.ExecutorService; +import java.util.concurrent.Executors; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.locks.Lock; +import java.util.concurrent.locks.ReentrantReadWriteLock; +import java.util.function.Consumer; +import javax.annotation.concurrent.ThreadSafe; +import lombok.SneakyThrows; +import lombok.extern.slf4j.Slf4j; +import org.apache.kafka.common.utils.Time; + +/** + * A system timer implementation. + */ +@Slf4j +@ThreadSafe +public class SystemTimer implements Timer { + + /** + * Create a system timer builder. + * + * @return a system timer builder. + */ + public static Builder builder() { + return new Builder(); + } + + /** + * Builder to build a system timer. + */ + public static class Builder { + + private String executorName; + private long tickMs = 1; + private int wheelSize = 20; + private long startMs = Time.SYSTEM.hiResClockMs(); + + private Builder() {} + + public Builder executorName(String executorName) { + this.executorName = executorName; + return this; + } + + public Builder tickMs(long tickMs) { + this.tickMs = tickMs; + return this; + } + + public Builder wheelSize(int wheelSize) { + this.wheelSize = wheelSize; + return this; + } + + public Builder startMs(long startMs) { + this.startMs = startMs; + return this; + } + + public SystemTimer build() { + Objects.requireNonNull(executorName, "No executor name is provided"); + + return new SystemTimer( + executorName, + tickMs, + wheelSize, + startMs + ); + } + + } + + private final ExecutorService taskExecutor; + private final DelayQueue delayQueue; + private final AtomicInteger taskCounter; + private final TimingWheel timingWheel; + + // Locks used to protect data structures while ticking + private final ReentrantReadWriteLock readWriteLock; + private final Lock readLock; + private final Lock writeLock; + private final Consumer reinsert; + + private SystemTimer(String executorName, + long tickMs, + int wheelSize, + long startMs) { + this.taskExecutor = Executors.newFixedThreadPool( + 1, new ThreadFactoryBuilder() + .setDaemon(false) + .setNameFormat("system-timer-%d") + .build() + ); + this.delayQueue = new DelayQueue(); + this.taskCounter = new AtomicInteger(0); + this.timingWheel = new TimingWheel( + tickMs, + wheelSize, + startMs, + taskCounter, + delayQueue + ); + this.readWriteLock = new ReentrantReadWriteLock(); + this.readLock = readWriteLock.readLock(); + this.writeLock = readWriteLock.writeLock(); + this.reinsert = timerTaskEntry -> addTimerTaskEntry(timerTaskEntry); + } + + @Override + public void add(TimerTask timerTask) { + readLock.lock(); + try { + addTimerTaskEntry(new TimerTaskEntry( + timerTask, timerTask.delayMs + Time.SYSTEM.hiResClockMs() + )); + } finally { + readLock.unlock(); + } + } + + private void addTimerTaskEntry(TimerTaskEntry timerTaskEntry) { + if (!timingWheel.add(timerTaskEntry)) { + // Already expired or cancelled + if (!timerTaskEntry.cancelled()) { + taskExecutor.submit(timerTaskEntry.timerTask()); + } + } + } + + @SneakyThrows + @Override + public boolean advanceClock(long timeoutMs) { + TimerTaskList bucket = delayQueue.poll(timeoutMs, TimeUnit.MILLISECONDS); + if (null != bucket) { + writeLock.lock(); + try { + while (null != bucket) { + timingWheel.advanceClock(bucket.getExpiration()); + bucket.flush(reinsert); + bucket = delayQueue.poll(); + } + } finally { + writeLock.unlock(); + } + return true; + } else { + return false; + } + } + + @Override + public int size() { + return taskCounter.get(); + } + + @Override + public void shutdown() { + taskExecutor.shutdown(); + } + +} diff --git a/src/main/java/io/streamnative/kop/utils/timer/Timer.java b/src/main/java/io/streamnative/kop/utils/timer/Timer.java new file mode 100644 index 0000000000..8f17321449 --- /dev/null +++ b/src/main/java/io/streamnative/kop/utils/timer/Timer.java @@ -0,0 +1,51 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.utils.timer; + +/** + * The timer interface to execute delayed operations. + */ +public interface Timer { + + /** + * Add a new task to this executor. It will be executed after the task's delay + * (beginning from the time of submission) + * + * @param timerTask the task to add + */ + void add(TimerTask timerTask); + + /** + * Advance the internal clock, executing any tasks whose expiration has been + * reached within the duration of the passed timeout. + * + * @param timeoutMs + * @return whether or not any tasks were executed + */ + boolean advanceClock(long timeoutMs); + + /** + * Get the number of tasks pending execution. + * + * @return the number of tasks + */ + int size(); + + /** + * Shutdown the timer service, leaving pending tasks unexecuted. + */ + void shutdown(); + + +} diff --git a/src/main/java/io/streamnative/kop/utils/timer/TimerTask.java b/src/main/java/io/streamnative/kop/utils/timer/TimerTask.java new file mode 100644 index 0000000000..0adbc766be --- /dev/null +++ b/src/main/java/io/streamnative/kop/utils/timer/TimerTask.java @@ -0,0 +1,50 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.utils.timer; + +import io.streamnative.kop.utils.timer.TimerTaskList.TimerTaskEntry; + +/** + * Timer task. + */ +public abstract class TimerTask implements Runnable { + + protected final long delayMs; + private TimerTaskEntry timerTaskEntry = null; + + protected TimerTask(long delayMs) { + this.delayMs = delayMs; + } + + public synchronized void cancel() { + if (null != timerTaskEntry) { + timerTaskEntry.remove(); + timerTaskEntry = null; + } + } + + synchronized void setTimerTaskEntry(TimerTaskEntry entry) { + // if this timerTask is already held by an existing timer task entry, + // we will remove such an entry first. + if (null != timerTaskEntry && timerTaskEntry != entry) { + timerTaskEntry.remove(); + } + timerTaskEntry = entry; + } + + synchronized TimerTaskEntry getTimerTaskEntry() { + return timerTaskEntry; + } + +} diff --git a/src/main/java/io/streamnative/kop/utils/timer/TimerTaskList.java b/src/main/java/io/streamnative/kop/utils/timer/TimerTaskList.java new file mode 100644 index 0000000000..1f6451883d --- /dev/null +++ b/src/main/java/io/streamnative/kop/utils/timer/TimerTaskList.java @@ -0,0 +1,209 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.utils.timer; + +import edu.umd.cs.findbugs.annotations.SuppressFBWarnings; +import java.util.concurrent.Delayed; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.atomic.AtomicLong; +import java.util.function.Consumer; +import javax.annotation.concurrent.ThreadSafe; +import lombok.Getter; +import lombok.experimental.Accessors; +import lombok.extern.slf4j.Slf4j; +import org.apache.kafka.common.utils.Time; + +/** + * The timer task list is a java implementation of Kafka implementation. + */ +@SuppressFBWarnings({ + "EQ_COMPARETO_USE_OBJECT_EQUALS", + "HE_EQUALS_USE_HASHCODE" +}) +@Slf4j +@ThreadSafe +public class TimerTaskList implements Delayed { + + private final AtomicInteger taskCounter; + private final AtomicLong expiration; + + // TimerTaskList forms a doubly linked cyclic list using a dummy root entry + // root.next points to the head + // root.prev points to the tail + private final TimerTaskEntry root; + + public TimerTaskList(AtomicInteger taskCounter) { + this.taskCounter = taskCounter; + this.root = new TimerTaskEntry(null, -1); + this.root.next = root; + this.root.prev = root; + this.expiration = new AtomicLong(-1L); + } + + // Set the bucket's expiration time + // Returns true if the expiration time is changed + public boolean setExpiration(long expirationMs) { + return expiration.getAndSet(expirationMs) != expirationMs; + } + + // Get the bucket's expiration time + public long getExpiration() { + return expiration.get(); + } + + public synchronized void forEach(Consumer f) { + TimerTaskEntry entry = root.next; + while (entry != root) { + final TimerTaskEntry nextEntry = entry.next; + if (!entry.cancelled()) { + f.accept(entry.timerTask); + } + entry = nextEntry; + } + } + + // add a timer task entry to this list + public void add(TimerTaskEntry timerTaskEntry) { + boolean done = false; + while (!done) { + // Remove the timer task entry if it is already in any other list + // We do this outside of the sync block below to avoid deadlocking. + // We may retry until timerTaskEntry.list becomes null. + timerTaskEntry.remove(); + + synchronized (this) { + synchronized (timerTaskEntry) { + if (timerTaskEntry.list == null) { + // put the timer task entry to the end of the list. (root.prev points to the tail entry) + TimerTaskEntry tail = root.prev; + timerTaskEntry.next = root; + timerTaskEntry.prev = tail; + timerTaskEntry.list = this; + tail.next = timerTaskEntry; + root.prev = timerTaskEntry; + taskCounter.incrementAndGet(); + done = true; + } + } + } + } + } + + // Remove the specified timer task entry from this list + public void remove(TimerTaskEntry timerTaskEntry) { + synchronized (this) { + synchronized (timerTaskEntry) { + if (timerTaskEntry.list == this) { + timerTaskEntry.next.prev = timerTaskEntry.prev; + timerTaskEntry.prev.next = timerTaskEntry.next; + timerTaskEntry.next = null; + timerTaskEntry.prev = null; + timerTaskEntry.list = null; + taskCounter.decrementAndGet(); + } + } + } + } + + // Remove all task entries and apply the supplied function to each of them + public synchronized void flush(Consumer f) { + TimerTaskEntry head = root.next; + while (head != root) { + remove(head); + f.accept(head); + head = root.next; + } + expiration.set(-1L); + } + + public long getDelay(TimeUnit unit) { + return unit.convert(Math.max(getExpiration() - Time.SYSTEM.hiResClockMs(), 0), TimeUnit.MILLISECONDS); + } + + @Override + public int compareTo(Delayed o) { + TimerTaskList other = (TimerTaskList) o; + + if (getExpiration() < other.getExpiration()) { + return -1; + } else if (getExpiration() > other.getExpiration()) { + return 1; + } else { + return 0; + } + } + + /** + * A timer task entry in the timer task list. + */ + @Accessors(fluent = true) + protected static class TimerTaskEntry implements Comparable { + + @Getter + private final TimerTask timerTask; + @Getter + private final long expirationMs; + private volatile TimerTaskList list = null; + private TimerTaskEntry next = null; + private TimerTaskEntry prev = null; + + public TimerTaskEntry(TimerTask timerTask, + long expirationMs) { + this.timerTask = timerTask; + this.expirationMs = expirationMs; + // if this timerTask is already held by an existing timer task entry, + // setTimerTaskEntry will remove it. + if (null != timerTask) { + timerTask.setTimerTaskEntry(this); + } + } + + public boolean cancelled() { + return timerTask.getTimerTaskEntry() != this; + } + + public void remove() { + TimerTaskList currentList = list; + // If remove is called when another thread is moving the entry from a task entry list to another, + // this may fail to remove the entry due to the change of value of list. Thus, we retry until the + // list becomes null. In a rare case, this thread sees null and exits the loop, but the other thread + // insert the entry to another list later. + while (currentList != null) { + currentList.remove(this); + currentList = list; + } + } + + @Override + public int compareTo(TimerTaskEntry o) { + return Long.compare(this.expirationMs, o.expirationMs); + } + + @Override + public boolean equals(Object obj) { + if (!(obj instanceof TimerTaskEntry)) { + return false; + } + TimerTaskEntry other = (TimerTaskEntry) obj; + return compareTo(other) == 0 + && list == other.list + && next == other.next + && prev == other.prev + && timerTask == other.timerTask; + } + } + + +} diff --git a/src/main/java/io/streamnative/kop/utils/timer/TimingWheel.java b/src/main/java/io/streamnative/kop/utils/timer/TimingWheel.java new file mode 100644 index 0000000000..051d66544b --- /dev/null +++ b/src/main/java/io/streamnative/kop/utils/timer/TimingWheel.java @@ -0,0 +1,206 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.utils.timer; + +import io.streamnative.kop.utils.timer.TimerTaskList.TimerTaskEntry; +import java.util.List; +import java.util.concurrent.DelayQueue; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Collectors; +import java.util.stream.IntStream; + +/** + * Hierarchical Timing Wheels + * + *

A simple timing wheel is a circular list of buckets of timer tasks. Let u be the time unit. + * A timing wheel with size n has n buckets and can hold timer tasks in n * u time interval. + * Each bucket holds timer tasks that fall into the corresponding time range. At the beginning, + * the first bucket holds tasks for [0, u), the second bucket holds tasks for [u, 2u), …, + * the n-th bucket for [u * (n -1), u * n). Every interval of time unit u, the timer ticks and + * moved to the next bucket then expire all timer tasks in it. So, the timer never insert a task + * into the bucket for the current time since it is already expired. The timer immediately runs + * the expired task. The emptied bucket is then available for the next round, so if the current + * bucket is for the time t, it becomes the bucket for [t + u * n, t + (n + 1) * u) after a tick. + * A timing wheel has O(1) cost for insert/delete (start-timer/stop-timer) whereas priority queue + * based timers, such as java.util.concurrent.DelayQueue and java.util.Timer, have O(log n) + * insert/delete cost. + * + *

A major drawback of a simple timing wheel is that it assumes that a timer request is within + * the time interval of n * u from the current time. If a timer request is out of this interval, + * it is an overflow. A hierarchical timing wheel deals with such overflows. It is a hierarchically + * organized timing wheels. The lowest level has the finest time resolution. As moving up the + * hierarchy, time resolutions become coarser. If the resolution of a wheel at one level is u and + * the size is n, the resolution of the next level should be n * u. At each level overflows are + * delegated to the wheel in one level higher. When the wheel in the higher level ticks, it reinsert + * timer tasks to the lower level. An overflow wheel can be created on-demand. When a bucket in an + * overflow bucket expires, all tasks in it are reinserted into the timer recursively. The tasks + * are then moved to the finer grain wheels or be executed. The insert (start-timer) cost is O(m) + * where m is the number of wheels, which is usually very small compared to the number of requests + * in the system, and the delete (stop-timer) cost is still O(1). + * + *

Example + * Let's say that u is 1 and n is 3. If the start time is c, + * then the buckets at different levels are: + * + *

+ * level buckets + * 1 [c,c] [c+1,c+1] [c+2,c+2] + * 2 [c,c+2] [c+3,c+5] [c+6,c+8] + * 3 [c,c+8] [c+9,c+17] [c+18,c+26] + *

+ * + *

The bucket expiration is at the time of bucket beginning. + * So at time = c+1, buckets [c,c], [c,c+2] and [c,c+8] are expired. + * Level 1's clock moves to c+1, and [c+3,c+3] is created. + * Level 2 and level3's clock stay at c since their clocks move in unit of 3 and 9, respectively. + * So, no new buckets are created in level 2 and 3. + * + *

Note that bucket [c,c+2] in level 2 won't receive any task since that range is already covered in level 1. + * The same is true for the bucket [c,c+8] in level 3 since its range is covered in level 2. + * This is a bit wasteful, but simplifies the implementation. + * + *

+ * 1 [c+1,c+1] [c+2,c+2] [c+3,c+3] + * 2 [c,c+2] [c+3,c+5] [c+6,c+8] + * 3 [c,c+8] [c+9,c+17] [c+18,c+26] + *

+ * + *

At time = c+2, [c+1,c+1] is newly expired. + * Level 1 moves to c+2, and [c+4,c+4] is created, + * + *

+ * 1 [c+2,c+2] [c+3,c+3] [c+4,c+4] + * 2 [c,c+2] [c+3,c+5] [c+6,c+8] + * 3 [c,c+8] [c+9,c+17] [c+18,c+18] + *

+ * + *

+ * At time = c+3, [c+2,c+2] is newly expired. + * Level 2 moves to c+3, and [c+5,c+5] and [c+9,c+11] are created. + * Level 3 stay at c. + *

+ * + *

+ * 1 [c+3,c+3] [c+4,c+4] [c+5,c+5] + * 2 [c+3,c+5] [c+6,c+8] [c+9,c+11] + * 3 [c,c+8] [c+9,c+17] [c+8,c+11] + *

+ * + *

The hierarchical timing wheels works especially well when operations are completed before they time out. + * Even when everything times out, it still has advantageous when there are many items in the timer. + * Its insert cost (including reinsert) and delete cost are O(m) and O(1), respectively while priority + * queue based timers takes O(log N) for both insert and delete where N is the number of items in the queue. + * + *

This class is not thread-safe. There should not be any add calls while advanceClock is executing. + * It is caller's responsibility to enforce it. Simultaneous add calls are thread-safe. + * + *

Note: this is the implementation from Kafka. + */ +class TimingWheel { + + private final long tickMs; + private final int wheelSize; + private final long startMs; + private final AtomicInteger taskCounter; + private final DelayQueue queue; + + private final long interval; + private final List buckets; + private long currentTime; + + // overflowWheel can potentially be updated and read by two concurrent threads through add(). + // Therefore, it needs to be volatile due to the issue of Double-Checked Locking pattern with JVM + private volatile TimingWheel overflowWheel = null; + + public TimingWheel( + long tickMs, + int wheelSize, + long startMs, + AtomicInteger taskCounter, + DelayQueue queue + ) { + this.tickMs = tickMs; + this.wheelSize = wheelSize; + this.startMs = startMs; + this.taskCounter = taskCounter; + this.queue = queue; + + this.interval = tickMs * wheelSize; + this.buckets = IntStream.range(0, wheelSize) + .mapToObj(i -> new TimerTaskList(taskCounter)) + .collect(Collectors.toList()); + this.currentTime = startMs - (startMs % tickMs); // rounding down to multiple of tickMs + } + + private synchronized void addOverflowWheel() { + if (null == overflowWheel) { + overflowWheel = new TimingWheel( + interval, + wheelSize, + currentTime, + taskCounter, + queue + ); + } + } + + public boolean add(TimerTaskEntry timerTaskEntry) { + final long expiration = timerTaskEntry.expirationMs(); + + if (timerTaskEntry.cancelled()) { + // cancelled + return false; + } else if (expiration < currentTime + tickMs) { + // Already expired + return false; + } else if (expiration < currentTime + interval) { + // Put in its own bucket + final long virtualId = expiration / tickMs; + TimerTaskList bucket = buckets.get( + (int) (virtualId % (long) wheelSize) + ); + bucket.add(timerTaskEntry); + + // Set the bucket expiration time + if (bucket.setExpiration(virtualId * tickMs)) { + // The bucket needs to be enqueued because it was an expired bucket + // We only need to enqueue the bucket when its expiration time has changed, i.e. the wheel has advanced + // and the previous buckets gets reused; further calls to set the expiration within the same wheel cycle + // will pass in the same value and hence return false, thus the bucket with the same expiration will not + // be enqueued multiple times. + queue.offer(bucket); + } + return true; + } else { + // Out of the interval. Put it into the parent timer + if (null == overflowWheel) { + addOverflowWheel(); + } + return overflowWheel.add(timerTaskEntry); + } + } + + // Try to advance the clock + public void advanceClock(long timeMs) { + if (timeMs >= currentTime + tickMs) { + currentTime = timeMs - (timeMs % tickMs); + + // Try to advance the clock of the overflow wheel if present + if (null != overflowWheel) { + overflowWheel.advanceClock(currentTime); + } + } + } + +} diff --git a/src/main/java/io/streamnative/kop/utils/timer/package-info.java b/src/main/java/io/streamnative/kop/utils/timer/package-info.java new file mode 100644 index 0000000000..ae50afe3b5 --- /dev/null +++ b/src/main/java/io/streamnative/kop/utils/timer/package-info.java @@ -0,0 +1,19 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +/** + * Timer related classes. + * + *

The classes under this package are ported from Kafka. + */ +package io.streamnative.kop.utils.timer; diff --git a/src/test/java/io/streamnative/kop/utils/MockTime.java b/src/test/java/io/streamnative/kop/utils/MockTime.java new file mode 100644 index 0000000000..ffc3086a9a --- /dev/null +++ b/src/test/java/io/streamnative/kop/utils/MockTime.java @@ -0,0 +1,111 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.utils; + +import java.util.concurrent.CopyOnWriteArrayList; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicLong; +import org.apache.kafka.common.utils.Time; + +/** + * A clock that you can manually advance by calling sleep. + */ +public class MockTime implements Time { + + /** + * Mock time listener. + */ + interface MockTimeListener { + void tick(); + } + + /** + * Listeners which are waiting for time changes. + */ + private final CopyOnWriteArrayList listeners = new CopyOnWriteArrayList<>(); + + private final long autoTickMs; + + // Values from `nanoTime` and `currentTimeMillis` are not comparable, so we store them separately to allow tests + // using this class to detect bugs where this is incorrectly assumed to be true + private final AtomicLong timeMs; + private final AtomicLong highResTimeNs; + + public MockTime() { + this(0); + } + + public MockTime(long autoTickMs) { + this(autoTickMs, System.currentTimeMillis(), System.nanoTime()); + } + + public MockTime(long autoTickMs, long currentTimeMs, long currentHighResTimeNs) { + this.timeMs = new AtomicLong(currentTimeMs); + this.highResTimeNs = new AtomicLong(currentHighResTimeNs); + this.autoTickMs = autoTickMs; + } + + public void addListener(MockTimeListener listener) { + listeners.add(listener); + } + + @Override + public long milliseconds() { + maybeSleep(autoTickMs); + return timeMs.get(); + } + + @Override + public long nanoseconds() { + maybeSleep(autoTickMs); + return highResTimeNs.get(); + } + + @Override + public long hiResClockMs() { + return TimeUnit.NANOSECONDS.toMillis(nanoseconds()); + } + + private void maybeSleep(long ms) { + if (ms != 0) { + sleep(ms); + } + } + + @Override + public void sleep(long ms) { + timeMs.addAndGet(ms); + highResTimeNs.addAndGet(TimeUnit.MILLISECONDS.toNanos(ms)); + tick(); + } + + public void setCurrentTimeMs(long newMs) { + long oldMs = timeMs.getAndSet(newMs); + + // does not allow to set to an older timestamp + if (oldMs > newMs) { + throw new IllegalArgumentException("Setting the time to " + newMs + " while current time " + + oldMs + " is newer; this is not allowed"); + } + + highResTimeNs.set(TimeUnit.MILLISECONDS.toNanos(newMs)); + tick(); + } + + private void tick() { + for (MockTimeListener listener : listeners) { + listener.tick(); + } + } +} diff --git a/src/test/java/io/streamnative/kop/utils/timer/MockTimer.java b/src/test/java/io/streamnative/kop/utils/timer/MockTimer.java new file mode 100644 index 0000000000..27d6b2c0f5 --- /dev/null +++ b/src/test/java/io/streamnative/kop/utils/timer/MockTimer.java @@ -0,0 +1,85 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.utils.timer; + +import io.streamnative.kop.utils.MockTime; +import io.streamnative.kop.utils.timer.TimerTaskList.TimerTaskEntry; +import java.util.Comparator; +import java.util.PriorityQueue; + +/** + * A mock implementation of {@link Timer}. + */ +public class MockTimer implements Timer { + + private final MockTime time = new MockTime(); + private final PriorityQueue taskQueue = new PriorityQueue<>(Comparator.reverseOrder()); + + @Override + public void add(TimerTask timerTask) { + if (timerTask.delayMs <= 0) { + timerTask.run(); + } else { + synchronized (taskQueue) { + taskQueue.add( + new TimerTaskEntry( + timerTask, timerTask.delayMs + time.milliseconds())); + } + } + } + + @Override + public boolean advanceClock(long timeoutMs) { + time.sleep(timeoutMs); + + boolean executed = false; + final long now = time.milliseconds(); + boolean hasMore = true; + + while (hasMore) { + hasMore = false; + TimerTaskEntry head; + synchronized (taskQueue) { + head = taskQueue.peek(); + if (null != head && now > head.expirationMs()) { + head = taskQueue.poll(); + hasMore = !taskQueue.isEmpty(); + } else { + head = null; + } + } + if (null != head) { + if (!head.cancelled()) { + TimerTask task = head.timerTask(); + task.run(); + executed = true; + } + } + } + + return executed; + } + + @Override + public int size() { + synchronized (taskQueue) { + return taskQueue.size(); + } + } + + @Override + public void shutdown() { + // no-op + } +} diff --git a/src/test/java/io/streamnative/kop/utils/timer/TimerTaskListTest.java b/src/test/java/io/streamnative/kop/utils/timer/TimerTaskListTest.java new file mode 100644 index 0000000000..89a512ad17 --- /dev/null +++ b/src/test/java/io/streamnative/kop/utils/timer/TimerTaskListTest.java @@ -0,0 +1,108 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.utils.timer; + +import static org.junit.Assert.assertEquals; + +import io.streamnative.kop.utils.timer.TimerTaskList.TimerTaskEntry; +import java.util.List; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import org.junit.Test; + +/** + * Unit test {@link TimerTaskList}. + */ +public class TimerTaskListTest { + + /** + * Test task. + */ + private static class TestTask extends TimerTask { + + protected TestTask(long delayMs) { + super(delayMs); + } + + @Override + public void run() { + + } + + } + + private int size(TimerTaskList list) { + AtomicInteger count = new AtomicInteger(0); + list.forEach(ignored -> count.incrementAndGet()); + return count.get(); + } + + @Test + public void testAll() { + AtomicInteger sharedCounter = new AtomicInteger(0); + TimerTaskList list1 = new TimerTaskList(sharedCounter); + TimerTaskList list2 = new TimerTaskList(sharedCounter); + TimerTaskList list3 = new TimerTaskList(sharedCounter); + + List tasks = IntStream.rangeClosed(1, 10).mapToObj(i -> { + TestTask task = new TestTask(0L); + list1.add(new TimerTaskEntry(task, 10L)); + assertEquals(i, sharedCounter.get()); + return task; + }).collect(Collectors.toList()); + + assertEquals(tasks.size(), sharedCounter.get()); + + // reinserting the existing tasks shouldn't change the task count. + tasks.subList(0, 4).forEach(task -> { + int prevCount = sharedCounter.get(); + // new TimerTaskEntry(task) will remove the existing entry from the list + list2.add(new TimerTaskEntry(task, 10L)); + assertEquals(prevCount, sharedCounter.get()); + }); + assertEquals(10 - 4, size(list1)); + assertEquals(4, size(list2)); + assertEquals(tasks.size(), sharedCounter.get()); + + // reinserting the existing tasks shouldn't change the task count + tasks.subList(4, 10).forEach(task -> { + int prevCount = sharedCounter.get(); + // new TimerTaskEntry(task) will remove the existing entry from the list + list3.add(new TimerTaskEntry(task, 10L)); + assertEquals(prevCount, sharedCounter.get()); + }); + assertEquals(0, size(list1)); + assertEquals(4, size(list2)); + assertEquals(6, size(list3)); + assertEquals(tasks.size(), sharedCounter.get()); + + // cancel tasks in the lists + list1.forEach(TimerTask::cancel); + assertEquals(0, size(list1)); + assertEquals(4, size(list2)); + assertEquals(6, size(list3)); + + list2.forEach(TimerTask::cancel); + assertEquals(0, size(list1)); + assertEquals(0, size(list2)); + assertEquals(6, size(list3)); + + list3.forEach(TimerTask::cancel); + assertEquals(0, size(list1)); + assertEquals(0, size(list2)); + assertEquals(0, size(list3)); + } + +} diff --git a/src/test/java/io/streamnative/kop/utils/timer/TimerTest.java b/src/test/java/io/streamnative/kop/utils/timer/TimerTest.java new file mode 100644 index 0000000000..5bbe356d11 --- /dev/null +++ b/src/test/java/io/streamnative/kop/utils/timer/TimerTest.java @@ -0,0 +1,161 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.utils.timer; + +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.fail; + +import com.google.common.collect.Lists; +import java.util.ArrayList; +import java.util.Collections; +import java.util.List; +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import lombok.extern.slf4j.Slf4j; +import org.apache.kafka.common.utils.Time; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +/** + * Unit test {@link Timer}. + */ +@Slf4j +public class TimerTest { + + private static class TestTask extends TimerTask { + private final int id; + private final CountDownLatch latch; + private final List output; + private final AtomicBoolean completed = new AtomicBoolean(false); + + public TestTask(long delayMs, + int id, + CountDownLatch latch, + List output) { + super(delayMs); + this.id = id; + this.latch = latch; + this.output = output; + } + + public void run() { + if (completed.compareAndSet(false, true)) { + synchronized (output) { + output.add(id); + } + latch.countDown(); + } + } + } + + private Timer timer = null; + + @Before + public void setup() { + this.timer = SystemTimer.builder() + .executorName("test") + .tickMs(1) + .wheelSize(3) + .startMs(Time.SYSTEM.hiResClockMs()) + .build(); + } + + @After + public void teardown() { + timer.shutdown(); + } + + @Test + public void testAlreadyExpiredTask() { + final List output = new ArrayList<>(); + final List latches = IntStream.range(-5, 0).mapToObj(i -> { + CountDownLatch latch = new CountDownLatch(1); + timer.add(new TestTask(i, i, latch, output)); + return latch; + }).collect(Collectors.toList()); + + timer.advanceClock(0); + + latches.forEach(latch -> { + try { + assertEquals( + "already expired tasks should run immediately", + true, + latch.await(3, TimeUnit.SECONDS)); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + fail("Should not reach here"); + } + }); + + assertEquals( + "Output of already expired tasks", + Lists.newArrayList(-5, -4, -3, -2, -1), + output + ); + } + + @Test + public void testTaskExpiration() { + final List output = new ArrayList<>(); + final List tasks = new ArrayList<>(); + final List ids = new ArrayList<>(); + final List latches = IntStream.range(0, 5).mapToObj(i -> { + CountDownLatch latch = new CountDownLatch(1); + tasks.add(new TestTask(i, i, latch, output)); + ids.add(i); + return latch; + }).collect(Collectors.toList()); + latches.addAll(IntStream.range(10, 100).mapToObj(i -> { + CountDownLatch latch = new CountDownLatch(1); + tasks.add(new TestTask(i, i, latch, output)); + tasks.add(new TestTask(i, i, latch, output)); + ids.add(i); + ids.add(i); + return latch; + }).collect(Collectors.toList())); + latches.addAll(IntStream.range(100, 500).mapToObj(i -> { + CountDownLatch latch = new CountDownLatch(1); + tasks.add(new TestTask(i, i, latch, output)); + ids.add(i); + return latch; + }).collect(Collectors.toList())); + + // randomly submit requests + tasks.forEach(task -> timer.add(task)); + + while (timer.advanceClock(2000)) {} + + latches.forEach(latch -> { + try { + latch.await(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + fail("Should not reach here"); + } + }); + + Collections.sort(ids); + assertEquals( + "output should match", + ids, + output + ); + } + +} From eb70d58ecb313a943b909b1ba2654a091687da5f Mon Sep 17 00:00:00 2001 From: Sijie Guo Date: Sat, 3 Aug 2019 21:48:08 +0800 Subject: [PATCH 5/8] Port delayed operations implementation from Kafka --- .../io/streamnative/kop/utils/CoreUtils.java | 44 ++ .../kop/utils/ShutdownableThread.java | 120 +++++ .../kop/utils/delayed/DelayedOperation.java | 147 ++++++ .../utils/delayed/DelayedOperationKey.java | 98 ++++ .../delayed/DelayedOperationPurgatory.java | 410 +++++++++++++++ .../kop/utils/delayed/package-info.java | 17 + .../io/streamnative/kop/utils/TestUtils.java | 43 ++ .../utils/delayed/DelayedOperationTest.java | 487 ++++++++++++++++++ 8 files changed, 1366 insertions(+) create mode 100644 src/main/java/io/streamnative/kop/utils/CoreUtils.java create mode 100644 src/main/java/io/streamnative/kop/utils/ShutdownableThread.java create mode 100644 src/main/java/io/streamnative/kop/utils/delayed/DelayedOperation.java create mode 100644 src/main/java/io/streamnative/kop/utils/delayed/DelayedOperationKey.java create mode 100644 src/main/java/io/streamnative/kop/utils/delayed/DelayedOperationPurgatory.java create mode 100644 src/main/java/io/streamnative/kop/utils/delayed/package-info.java create mode 100644 src/test/java/io/streamnative/kop/utils/TestUtils.java create mode 100644 src/test/java/io/streamnative/kop/utils/delayed/DelayedOperationTest.java diff --git a/src/main/java/io/streamnative/kop/utils/CoreUtils.java b/src/main/java/io/streamnative/kop/utils/CoreUtils.java new file mode 100644 index 0000000000..3987086b31 --- /dev/null +++ b/src/main/java/io/streamnative/kop/utils/CoreUtils.java @@ -0,0 +1,44 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.utils; + +import java.util.concurrent.locks.Lock; +import java.util.concurrent.locks.ReadWriteLock; +import java.util.function.Supplier; +import lombok.experimental.UtilityClass; + +/** + * Core utils. + */ +@UtilityClass +public final class CoreUtils { + + public static T inLock(Lock lock, Supplier supplier) { + lock.lock(); + try { + return supplier.get(); + } finally { + lock.unlock(); + } + } + + public static T inReadLock(ReadWriteLock lock, Supplier supplier) { + return inLock(lock.readLock(), supplier); + } + + public static T inWriteLock(ReadWriteLock lock, Supplier supplier) { + return inLock(lock.writeLock(), supplier); + } + +} diff --git a/src/main/java/io/streamnative/kop/utils/ShutdownableThread.java b/src/main/java/io/streamnative/kop/utils/ShutdownableThread.java new file mode 100644 index 0000000000..bc05f25c53 --- /dev/null +++ b/src/main/java/io/streamnative/kop/utils/ShutdownableThread.java @@ -0,0 +1,120 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.utils; + +import java.util.concurrent.CountDownLatch; +import java.util.concurrent.TimeUnit; +import lombok.extern.slf4j.Slf4j; +import org.apache.kafka.common.internals.FatalExitError; +import org.apache.kafka.common.utils.Exit; + +/** + * Shutdownable thread. + */ +@Slf4j +public abstract class ShutdownableThread extends Thread { + + private final boolean isInterruptible; + private final String logIdent; + private final CountDownLatch shutdownInitiated = new CountDownLatch(1); + private final CountDownLatch shutdownComplete = new CountDownLatch(1); + + public ShutdownableThread(String name) { + this(name, true); + } + + public ShutdownableThread(String name, + boolean isInterruptible) { + super(name); + this.isInterruptible = isInterruptible; + this.setDaemon(false); + this.logIdent = "[" + name + "]"; + } + + public boolean isRunning() { + return shutdownInitiated.getCount() != 0; + } + + public void shutdown() throws InterruptedException { + initiateShutdown(); + awaitShutdown(); + } + + public boolean isShutdownComplete() { + return shutdownComplete.getCount() == 0; + } + + public synchronized boolean initiateShutdown() { + if (isRunning()) { + log.info("{} Shutting down", logIdent); + } + shutdownInitiated.countDown(); + if (isInterruptible) { + interrupt(); + return true; + } else { + return false; + } + } + + /** + * After calling initiateShutdown(), use this API to wait until the shutdown is complete. + */ + public void awaitShutdown() throws InterruptedException { + shutdownComplete.await(); + log.info("{} Shutdown completed", logIdent); + } + + /** + * Causes the current thread to wait until the shutdown is initiated, + * or the specified waiting time elapses. + * + * @param timeout + * @param unit + */ + public void pause(long timeout, TimeUnit unit) throws InterruptedException { + if (shutdownInitiated.await(timeout, unit)) { + if (log.isTraceEnabled()) { + log.trace("{} shutdownInitiated latch count reached zero. Shutdown called.", logIdent); + } + } + } + + /** + * This method is repeatedly invoked until the thread shuts down or this method throws an exception. + */ + protected abstract void doWork(); + + @Override + public void run() { + log.info("{} Starting", logIdent); + try { + while (isRunning()) { + doWork(); + } + } catch (FatalExitError e) { + shutdownInitiated.countDown(); + shutdownComplete.countDown(); + log.info("{} Stopped", logIdent); + Exit.exit(e.statusCode()); + } catch (Throwable cause) { + if (isRunning()) { + log.error("{} Error due to", logIdent, cause); + } + } finally { + shutdownComplete.countDown(); + } + } + +} diff --git a/src/main/java/io/streamnative/kop/utils/delayed/DelayedOperation.java b/src/main/java/io/streamnative/kop/utils/delayed/DelayedOperation.java new file mode 100644 index 0000000000..124afb7e35 --- /dev/null +++ b/src/main/java/io/streamnative/kop/utils/delayed/DelayedOperation.java @@ -0,0 +1,147 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.utils.delayed; + +import io.streamnative.kop.utils.timer.TimerTask; +import java.util.Optional; +import java.util.concurrent.atomic.AtomicBoolean; +import java.util.concurrent.locks.Lock; +import java.util.concurrent.locks.ReentrantLock; +import lombok.extern.slf4j.Slf4j; + +/** + * An operation whose processing needs to be delayed for at most the given delayMs. For example + * a delayed produce operation could be waiting for specified number of acks; or + * a delayed fetch operation could be waiting for a given number of bytes to accumulate. + * + *

The logic upon completing a delayed operation is defined in onComplete() and will be called exactly once. + * Once an operation is completed, isCompleted() will return true. onComplete() can be triggered by either + * forceComplete(), which forces calling onComplete() after delayMs if the operation is not yet completed, + * or tryComplete(), which first checks if the operation can be completed or not now, and if yes calls + * forceComplete(). + * + *

A subclass of DelayedOperation needs to provide an implementation of both onComplete() and tryComplete(). + */ +@Slf4j +public abstract class DelayedOperation extends TimerTask { + + protected final Optional lockOpt; + private final AtomicBoolean completed = new AtomicBoolean(false); + private final AtomicBoolean tryCompletePending = new AtomicBoolean(false); + final Lock lock; + + protected DelayedOperation(long delayMs, + Optional lockOpt) { + super(delayMs); + this.lockOpt = lockOpt; + this.lock = lockOpt.orElseGet(() -> new ReentrantLock()); + } + + /** + * Force completing the delayed operation, if not already completed. + * This function can be triggered when + * + *

1. The operation has been verified to be completable inside tryComplete() + * 2. The operation has expired and hence needs to be completed right now + * + *

Return true iff the operation is completed by the caller: note that + * concurrent threads can try to complete the same operation, but only + * the first thread will succeed in completing the operation and return + * true, others will still return false + */ + public boolean forceComplete() { + if (completed.compareAndSet(false, true)) { + // cancel the timeout timer + cancel(); + onComplete(); + return true; + } else { + return false; + } + } + + /** + * Check if the delayed operation is already completed. + */ + public boolean isCompleted() { + return completed.get(); + } + + /** + * Call-back to execute when a delayed operation gets expired and hence forced to complete. + */ + public abstract void onExpiration(); + + /** + * Process for completing an operation. This function needs to be defined + * in subclasses and will be called exactly once in forceComplete() + */ + public abstract void onComplete(); + + /** + * Try to complete the delayed operation by first checking if the operation + * can be completed by now. If yes execute the completion logic by calling + * forceComplete() and return true iff forceComplete returns true; otherwise return false + * + *

This function needs to be defined in subclasses. + */ + public abstract boolean tryComplete(); + + /** + * Thread-safe variant of tryComplete() that attempts completion only if the lock can be acquired + * without blocking. + * + *

If threadA acquires the lock and performs the check for completion before completion criteria is met + * and threadB satisfies the completion criteria, but fails to acquire the lock because threadA has not + * yet released the lock, we need to ensure that completion is attempted again without blocking threadA + * or threadB. `tryCompletePending` is set by threadB when it fails to acquire the lock and at least one + * of threadA or threadB will attempt completion of the operation if this flag is set. This ensures that + * every invocation of `maybeTryComplete` is followed by at least one invocation of `tryComplete` until + * the operation is actually completed. + */ + boolean maybeTryComplete() { + boolean retry = false; + boolean done = false; + do { + if (lock.tryLock()) { + try { + tryCompletePending.set(false); + done = tryComplete(); + } finally { + lock.unlock(); + } + // While we were holding the lock, another thread may have invoked `maybeTryComplete` and set + // `tryCompletePending`. In this case we should retry. + retry = tryCompletePending.get(); + } else { + // Another thread is holding the lock. If `tryCompletePending` is already set and this thread failed to + // acquire the lock, then the thread that is holding the lock is guaranteed to see the flag and retry. + // Otherwise, we should set the flag and retry on this thread since the thread holding the lock may have + // released the lock and returned by the time the flag is set. + retry = !tryCompletePending.getAndSet(true); + } + } while (!isCompleted() && retry); + return done; + } + + /** + * run() method defines a task that is executed on timeout. + */ + @Override + public void run() { + if (forceComplete()) { + onExpiration(); + } + } +} diff --git a/src/main/java/io/streamnative/kop/utils/delayed/DelayedOperationKey.java b/src/main/java/io/streamnative/kop/utils/delayed/DelayedOperationKey.java new file mode 100644 index 0000000000..03e7839560 --- /dev/null +++ b/src/main/java/io/streamnative/kop/utils/delayed/DelayedOperationKey.java @@ -0,0 +1,98 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.utils.delayed; + +import lombok.Data; +import lombok.RequiredArgsConstructor; +import lombok.experimental.Accessors; + +/** + * Delayed operation key. + */ +public interface DelayedOperationKey { + + /** + * Key label. + * + * @return key label. + */ + String keyLabel(); + + /** + * Member key. + */ + @Data + @Accessors(fluent = true) + @RequiredArgsConstructor + class MemberKey implements DelayedOperationKey { + + private final String groupId; + private final String consumerId; + + @Override + public String keyLabel() { + return String.format("%s-%s", groupId, consumerId); + } + } + + /** + * Group key. + */ + @Data + @Accessors(fluent = true) + @RequiredArgsConstructor + class GroupKey implements DelayedOperationKey { + + private final String groupId; + + @Override + public String keyLabel() { + return groupId; + } + } + + /** + * Topic key. + */ + @Data + @Accessors(fluent = true) + @RequiredArgsConstructor + class TopicKey implements DelayedOperationKey { + + private final String topic; + + @Override + public String keyLabel() { + return topic; + } + } + + /** + * Topic partition key. + */ + @Data + @Accessors(fluent = true) + @RequiredArgsConstructor + class TopicPartitionOperationKey implements DelayedOperationKey { + + private final String topic; + private final int partition; + + @Override + public String keyLabel() { + return String.format("%s-%d", topic, partition); + } + } + +} diff --git a/src/main/java/io/streamnative/kop/utils/delayed/DelayedOperationPurgatory.java b/src/main/java/io/streamnative/kop/utils/delayed/DelayedOperationPurgatory.java new file mode 100644 index 0000000000..89170e315c --- /dev/null +++ b/src/main/java/io/streamnative/kop/utils/delayed/DelayedOperationPurgatory.java @@ -0,0 +1,410 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.utils.delayed; + +import static com.google.common.base.Preconditions.checkArgument; +import static io.streamnative.kop.utils.CoreUtils.inReadLock; +import static io.streamnative.kop.utils.CoreUtils.inWriteLock; + +import io.streamnative.kop.utils.ShutdownableThread; +import io.streamnative.kop.utils.timer.SystemTimer; +import io.streamnative.kop.utils.timer.Timer; +import java.util.ArrayList; +import java.util.Collection; +import java.util.Collections; +import java.util.Iterator; +import java.util.List; +import java.util.concurrent.ConcurrentHashMap; +import java.util.concurrent.ConcurrentLinkedQueue; +import java.util.concurrent.ConcurrentMap; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.locks.ReentrantReadWriteLock; +import lombok.extern.slf4j.Slf4j; + +/** + * A helper purgatory class for bookkeeping delayed operations with a timeout, and expiring timed out operations. + */ +@Slf4j +public class DelayedOperationPurgatory { + + public static Builder builder() { + return new Builder<>(); + } + + /** + * Builder to build a delayed operation purgatory. + */ + public static class Builder { + + private String purgatoryName; + private Timer timer; + private int purgeInterval = 1000; + private boolean reaperEnabled = true; + private boolean timerEnabled = true; + + private Builder() {} + + public Builder purgatoryName(String purgatoryName) { + this.purgatoryName = purgatoryName; + return this; + } + + public Builder timeoutTimer(Timer timer) { + this.timer = timer; + return this; + } + + public Builder purgeInterval(int purgeInterval) { + this.purgeInterval = purgeInterval; + return this; + } + + public Builder reaperEnabled(boolean reaperEnabled) { + this.reaperEnabled = reaperEnabled; + return this; + } + + public Builder timerEnabled(boolean timerEnabled) { + this.timerEnabled = timerEnabled; + return this; + } + + public DelayedOperationPurgatory build() { + if (null == timer) { + timer = SystemTimer.builder().executorName(purgatoryName).build(); + } + return new DelayedOperationPurgatory<>( + purgatoryName, + timer, + purgeInterval, + reaperEnabled, + timerEnabled + ); + } + } + + private final String purgatoryName; + private final Timer timeoutTimer; + private final int purgeInterval; + private final boolean reaperEnabled; + private final boolean timerEnabled; + + /* a list of operation watching keys */ + private final ConcurrentMap watchersForKey; + + private final ReentrantReadWriteLock removeWatchersLock = new ReentrantReadWriteLock(); + + // the number of estimated total operations in the purgatory + private final AtomicInteger estimatedTotalOperations = new AtomicInteger(0); + + /* background thread expiring operations that have timed out */ + private final ShutdownableThread expirationReaper; + + public DelayedOperationPurgatory( + String purgatoryName, + Timer timeoutTimer, + int purgeInterval, + boolean reaperEnabled, + boolean timerEnabled + ) { + this.purgatoryName = purgatoryName; + this.timeoutTimer = timeoutTimer; + this.purgeInterval = purgeInterval; + this.reaperEnabled = reaperEnabled; + this.timerEnabled = timerEnabled; + + this.watchersForKey = new ConcurrentHashMap<>(); + this.expirationReaper = new ShutdownableThread( + String.format("ExpirationReaper-%s", purgatoryName) + ) { + @Override + protected void doWork() { + advanceClock(200L); + } + }; + + if (reaperEnabled) { + expirationReaper.start(); + } + } + + /** + * Check if the operation can be completed, if not watch it based on the given watch keys + * + *

Note that a delayed operation can be watched on multiple keys. It is possible that + * an operation is completed after it has been added to the watch list for some, but + * not all of the keys. In this case, the operation is considered completed and won't + * be added to the watch list of the remaining keys. The expiration reaper thread will + * remove this operation from any watcher list in which the operation exists. + * + * @param operation the delayed operation to be checked + * @param watchKeys keys for bookkeeping the operation + * @return true iff the delayed operations can be completed by the caller + */ + public boolean tryCompleteElseWatch(T operation, List watchKeys) { + checkArgument(!watchKeys.isEmpty(), "The watch key list can't be empty"); + + // The cost of tryComplete() is typically proportional to the number of keys. Calling + // tryComplete() for each key is going to be expensive if there are many keys. Instead, + // we do the check in the following way. Call tryComplete(). If the operation is not completed, + // we just add the operation to all keys. Then we call tryComplete() again. At this time, if + // the operation is still not completed, we are guaranteed that it won't miss any future triggering + // event since the operation is already on the watcher list for all keys. This does mean that + // if the operation is completed (by another thread) between the two tryComplete() calls, the + // operation is unnecessarily added for watch. However, this is a less severe issue since the + // expire reaper will clean it up periodically. + + // At this point the only thread that can attempt this operation is this current thread + // Hence it is safe to tryComplete() without a lock + boolean isCompletedByMe = operation.tryComplete(); + if (isCompletedByMe) { + return true; + } + + boolean watchCreated = false; + for (Object key : watchKeys) { + // If the operation is already completed, stop adding it to the rest of the watcher list. + if (operation.isCompleted()) { + return false; + } + watchForOperation(key, operation); + + if (!watchCreated) { + watchCreated = true; + estimatedTotalOperations.incrementAndGet(); + } + } + + isCompletedByMe = operation.maybeTryComplete(); + if (isCompletedByMe) { + return true; + } + + // if it cannot be completed by now and hence is watched, add to the expire queue also + if (!operation.isCompleted()) { + if (timerEnabled) { + timeoutTimer.add(operation); + } + if (operation.isCompleted()) { + // cancel the timer task + operation.cancel(); + } + } + + return false; + } + + /** + * Check if some delayed operations can be completed with the given watch key, + * and if yes complete them. + * + * @return the number of completed operations during this process + */ + public int checkAndComplete(Object key) { + Watchers watchers = inReadLock( + removeWatchersLock, + () -> watchersForKey.get(key)); + if (null == watchers) { + return 0; + } else { + return watchers.tryCompleteWatched(); + } + } + + /** + * Return the total size of watch lists the purgatory. Since an operation may be watched + * on multiple lists, and some of its watched entries may still be in the watch lists + * even when it has been completed, this number may be larger than the number of real operations watched + */ + public int watched() { + return allWatchers().stream().mapToInt(Watchers::countWatched).sum(); + } + + /** + * Return the number of delayed operations in the expiry queue. + */ + public int delayed() { + return timeoutTimer.size(); + } + + /** + * Cancel watching on any delayed operations for the given key. Note the operation will not be completed + */ + public List cancelForKey(Object key) { + return inWriteLock(removeWatchersLock, () -> { + Watchers watchers = watchersForKey.remove(key); + if (watchers != null) { + return watchers.cancel(); + } else { + return Collections.emptyList(); + } + }); + } + /* + * Return all the current watcher lists, + * note that the returned watchers may be removed from the list by other threads + */ + private Collection allWatchers() { + return inReadLock(removeWatchersLock, () -> watchersForKey.values()); + } + + /* + * Return the watch list of the given key, note that we need to + * grab the removeWatchersLock to avoid the operation being added to a removed watcher list + */ + private void watchForOperation(Object key, T operation) { + inReadLock(removeWatchersLock, () -> { + watchersForKey.computeIfAbsent(key, (k) -> new Watchers(k)) + .watch(operation); + return null; + }); + } + + /** + * Remove the key from watcher lists if its list is empty. + */ + private void removeKeyIfEmpty(Object key, Watchers watchers) { + inWriteLock(removeWatchersLock, () -> { + // if the current key is no longer correlated to the watchers to remove, skip + if (watchersForKey.get(key) != watchers) { + return null; + } + + if (watchers != null && watchers.isEmpty()) { + watchersForKey.remove(key); + } + return null; + }); + } + + /** + * Shutdown the expire reaper thread. + */ + public void shutdown() { + if (reaperEnabled) { + try { + expirationReaper.shutdown(); + } catch (InterruptedException e) { + Thread.currentThread().interrupt(); + log.error("Interrupted at shutting down expiration reaper for {}", purgatoryName); + } + } + timeoutTimer.shutdown(); + } + + /** + * A linked list of watched delayed operations based on some key. + */ + private class Watchers { + + private final Object key; + private final ConcurrentLinkedQueue operations = new ConcurrentLinkedQueue<>(); + + Watchers(Object key) { + this.key = key; + } + + // count the current number of watched operations. This is O(n), so use isEmpty() if possible + public int countWatched() { + return operations.size(); + } + + public boolean isEmpty() { + return operations.isEmpty(); + } + + // add the element to watch + public void watch(T t) { + operations.add(t); + } + + // traverse the list and try to complete some watched elements + public int tryCompleteWatched() { + int completed = 0; + + Iterator iter = operations.iterator(); + while (iter.hasNext()) { + T curr = iter.next(); + if (curr.isCompleted()) { + // another thread has completed this operation, just remove it + iter.remove(); + } else if (curr.maybeTryComplete()) { + iter.remove(); + completed += 1; + } + } + + if (operations.isEmpty()) { + removeKeyIfEmpty(key, this); + } + + return completed; + } + + public List cancel() { + Iterator iter = operations.iterator(); + List cancelled = new ArrayList<>(); + while (iter.hasNext()) { + T curr = iter.next(); + curr.cancel(); + iter.remove(); + cancelled.add(curr); + } + return cancelled; + } + + // traverse the list and purge elements that are already completed by others + int purgeCompleted() { + int purged = 0; + + Iterator iter = operations.iterator(); + while (iter.hasNext()) { + T curr = iter.next(); + if (curr.isCompleted()) { + iter.remove(); + purged += 1; + } + } + + if (operations.isEmpty()) { + removeKeyIfEmpty(key, this); + } + + return purged; + } + } + + public void advanceClock(long timeoutMs) { + timeoutTimer.advanceClock(timeoutMs); + + // Trigger a purge if the number of completed but still being watched operations is larger than + // the purge threshold. That number is computed by the difference btw the estimated total number of + // operations and the number of pending delayed operations. + if (estimatedTotalOperations.get() - delayed() > purgeInterval) { + // now set estimatedTotalOperations to delayed (the number of pending operations) since we are going to + // clean up watchers. Note that, if more operations are completed during the clean up, we may end up with + // a little overestimated total number of operations. + estimatedTotalOperations.getAndSet(delayed()); + if (log.isDebugEnabled()) { + log.debug("{} Begin purging watch lists", purgatoryName); + } + int purged = allWatchers().stream().mapToInt(Watchers::purgeCompleted).sum(); + if (log.isDebugEnabled()) { + log.debug("{} Purged {} elements from watch lists.", purgatoryName, purged); + } + } + } + + +} diff --git a/src/main/java/io/streamnative/kop/utils/delayed/package-info.java b/src/main/java/io/streamnative/kop/utils/delayed/package-info.java new file mode 100644 index 0000000000..d8ec5f0e54 --- /dev/null +++ b/src/main/java/io/streamnative/kop/utils/delayed/package-info.java @@ -0,0 +1,17 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +/** + * Classes related delayed operations. + */ +package io.streamnative.kop.utils.delayed; \ No newline at end of file diff --git a/src/test/java/io/streamnative/kop/utils/TestUtils.java b/src/test/java/io/streamnative/kop/utils/TestUtils.java new file mode 100644 index 0000000000..ea2a19d1a9 --- /dev/null +++ b/src/test/java/io/streamnative/kop/utils/TestUtils.java @@ -0,0 +1,43 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.utils; + +import static org.junit.Assert.fail; + +import java.util.function.Supplier; +import lombok.SneakyThrows; + +/** + * Test utilities. + */ +public class TestUtils { + + @SneakyThrows + public static void waitUntilTrue(Supplier condition, + Supplier msg, + long waitTime, + long pause) { + long startTime = System.currentTimeMillis(); + while (true) { + if (condition.get()) { + return; + } + if (System.currentTimeMillis() > startTime + waitTime) { + fail(msg.get()); + } + Thread.sleep(Math.min(waitTime, pause)); + } + } + +} diff --git a/src/test/java/io/streamnative/kop/utils/delayed/DelayedOperationTest.java b/src/test/java/io/streamnative/kop/utils/delayed/DelayedOperationTest.java new file mode 100644 index 0000000000..0ef3eb15f9 --- /dev/null +++ b/src/test/java/io/streamnative/kop/utils/delayed/DelayedOperationTest.java @@ -0,0 +1,487 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.utils.delayed; + +import static io.streamnative.kop.utils.CoreUtils.inLock; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import com.google.common.collect.Lists; +import io.streamnative.kop.utils.TestUtils; +import java.util.List; +import java.util.Optional; +import java.util.Random; +import java.util.Set; +import java.util.concurrent.ExecutionException; +import java.util.concurrent.Executors; +import java.util.concurrent.Future; +import java.util.concurrent.ScheduledExecutorService; +import java.util.concurrent.Semaphore; +import java.util.concurrent.ThreadLocalRandom; +import java.util.concurrent.TimeUnit; +import java.util.concurrent.atomic.AtomicInteger; +import java.util.concurrent.locks.Lock; +import java.util.concurrent.locks.ReentrantLock; +import java.util.function.BiFunction; +import java.util.function.Function; +import java.util.function.Supplier; +import java.util.stream.Collectors; +import java.util.stream.IntStream; +import lombok.SneakyThrows; +import org.apache.kafka.common.utils.Time; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + +/** + * Unit test {@link DelayedOperation}. + */ +public class DelayedOperationTest { + + /** + * A mock delayed operation. + */ + static class MockDelayedOperation extends DelayedOperation { + + private final Optional responseLockOpt; + boolean completable = false; + + protected MockDelayedOperation(long delayMs) { + this(delayMs, Optional.empty(), Optional.empty()); + } + + protected MockDelayedOperation(long delayMs, + Optional lockOpt, + Optional responseLockOpt) { + super(delayMs, lockOpt); + this.responseLockOpt = responseLockOpt; + } + + public synchronized void awaitExpiration() throws InterruptedException { + wait(); + } + + @Override + public void onExpiration() { + // no-op + } + + @Override + public void onComplete() { + responseLockOpt.map(lock -> { + if (!lock.tryLock()) { + throw new IllegalStateException("Response callback lock could not be acquired in callback"); + } + return null; + }); + synchronized (this) { + notify(); + } + } + + @Override + public boolean tryComplete() { + if (completable) { + return forceComplete(); + } else { + return false; + } + } + } + + static class TestDelayOperation extends MockDelayedOperation { + + private final int index; + private final Object key; + private final AtomicInteger completionAttemptsRemaining; + private final int maxDelayMs; + + TestDelayOperation(int index, + int completionAttempts, + int maxDelayMs) { + super(10000L); + this.index = index; + this.key = "key" + index; + this.maxDelayMs = maxDelayMs; + this.completionAttemptsRemaining = new AtomicInteger(completionAttempts); + } + + @SneakyThrows + @Override + public boolean tryComplete() { + boolean shouldComplete = completable; + Thread.sleep(ThreadLocalRandom.current().nextInt(maxDelayMs)); + if (shouldComplete) { + return forceComplete(); + } else { + return false; + } + } + } + + DelayedOperationPurgatory purgatory = null; + ScheduledExecutorService executorService = null; + + @Before + public void setup() { + purgatory = DelayedOperationPurgatory.builder() + .purgatoryName("mock") + .build(); + } + + @After + public void teardown() { + purgatory.shutdown(); + if (null != executorService) { + executorService.shutdown(); + } + } + + @Test + public void testRequestSatisfaction() { + MockDelayedOperation r1 = new MockDelayedOperation(100000L); + MockDelayedOperation r2 = new MockDelayedOperation(100000L); + assertEquals( + "With no waiting requests, nothing should be satisfied", + 0, purgatory.checkAndComplete("test1")); + assertFalse( + "r1 not satisfied and hence watched", + purgatory.tryCompleteElseWatch(r1, Lists.newArrayList("test1"))); + assertEquals( + "Still nothing satisfied", + 0, purgatory.checkAndComplete("test1")); + assertFalse( + "r2 not satisfied and hence watched", + purgatory.tryCompleteElseWatch(r2, Lists.newArrayList("test2"))); + assertEquals( + "Still nothing satisfied", + 0, purgatory.checkAndComplete("test2")); + r1.completable = true; + assertEquals( + "r1 satisfied", + 1, purgatory.checkAndComplete("test1")); + assertEquals( + "Nothing satisfied", + 0, purgatory.checkAndComplete("test1")); + r2.completable = true; + assertEquals( + "r2 satisfied", + 1, purgatory.checkAndComplete("test2")); + assertEquals( + "Nothing satisfied", + 0, purgatory.checkAndComplete("test2")); + } + + @Test + public void testRequestExpiry() throws Exception { + long expiration = 20L; + long start = Time.SYSTEM.hiResClockMs(); + MockDelayedOperation r1 = new MockDelayedOperation(expiration); + MockDelayedOperation r2 = new MockDelayedOperation(200000L); + assertFalse( + "r1 not satisfied and hence watched", + purgatory.tryCompleteElseWatch(r1, Lists.newArrayList("test1"))); + assertFalse( + "r2 not satisfied and hence watched", + purgatory.tryCompleteElseWatch(r2, Lists.newArrayList("test2"))); + r1.awaitExpiration(); + long elapsed = Time.SYSTEM.hiResClockMs() - start; + assertTrue + ("r1 completed due to expiration", + r1.isCompleted()); + assertFalse("r2 hasn't completed", r2.isCompleted()); + assertTrue( + "Time for expiration $elapsed should at least " + expiration, + elapsed >= expiration); + } + + @Test + public void testRequestPurge() { + MockDelayedOperation r1 = new MockDelayedOperation(100000L); + MockDelayedOperation r2 = new MockDelayedOperation(100000L); + MockDelayedOperation r3 = new MockDelayedOperation(100000L); + purgatory.tryCompleteElseWatch(r1, Lists.newArrayList("test1")); + purgatory.tryCompleteElseWatch(r2, Lists.newArrayList("test1", "test2")); + purgatory.tryCompleteElseWatch(r3, Lists.newArrayList("test1", "test2", "test3")); + + assertEquals( + "Purgatory should have 3 total delayed operations", + 3, purgatory.delayed()); + assertEquals( + "Purgatory should have 6 watched elements", + 6, purgatory.watched()); + + // complete the operations, it should immediately be purged from the delayed operation + r2.completable = true; + r2.tryComplete(); + assertEquals( + "Purgatory should have 2 total delayed operations instead of " + purgatory.delayed(), + 2, purgatory.delayed()); + + r3.completable = true; + r3.tryComplete(); + assertEquals( + "Purgatory should have 1 total delayed operations instead of " + purgatory.delayed(), + 1, purgatory.delayed()); + + // checking a watch should purge the watch list + purgatory.checkAndComplete("test1"); + assertEquals( + "Purgatory should have 4 watched elements instead of " + purgatory.watched(), + 4, purgatory.watched()); + + purgatory.checkAndComplete("test2"); + assertEquals( + "Purgatory should have 2 watched elements instead of " + purgatory.watched(), + 2, purgatory.watched()); + + purgatory.checkAndComplete("test3"); + assertEquals( + "Purgatory should have 1 watched elements instead of " + purgatory.watched(), + 1, purgatory.watched()); + } + + @Test + public void shouldCancelForKeyReturningCancelledOperations() { + purgatory.tryCompleteElseWatch(new MockDelayedOperation(10000L), Lists.newArrayList("key")); + purgatory.tryCompleteElseWatch(new MockDelayedOperation(10000L), Lists.newArrayList("key")); + purgatory.tryCompleteElseWatch(new MockDelayedOperation(10000L), Lists.newArrayList("key2")); + + List cancelledOperations = purgatory.cancelForKey("key"); + assertEquals(2, cancelledOperations.size()); + assertEquals(1, purgatory.delayed()); + assertEquals(1, purgatory.watched()); + } + + @Test + public void shouldReturnNilOperationsOnCancelForKeyWhenKeyDoesntExist() { + List cancelledOperations = purgatory.cancelForKey("key"); + assertTrue(cancelledOperations.isEmpty()); + } + + /** + * Verify that if there is lock contention between two threads attempting to complete, + * completion is performed without any blocking in either thread. + */ + @Test + public void testTryCompleteLockContention() throws Exception { + executorService = Executors.newSingleThreadScheduledExecutor(); + AtomicInteger completionAttemptsRemaining = new AtomicInteger(Integer.MAX_VALUE); + Semaphore tryCompleteSemaphore = new Semaphore(1); + String key = "key"; + + MockDelayedOperation op = new MockDelayedOperation(100000L) { + @SneakyThrows + @Override + public boolean tryComplete() { + boolean shouldComplete = completionAttemptsRemaining.decrementAndGet() <= 0; + tryCompleteSemaphore.acquire(); + try { + if (shouldComplete) { + return forceComplete(); + } else { + return false; + } + } finally { + tryCompleteSemaphore.release(); + } + } + }; + + purgatory.tryCompleteElseWatch(op, Lists.newArrayList(key)); + completionAttemptsRemaining.set(2); + tryCompleteSemaphore.acquire(); + Future future = runOnAnotherThread(() -> purgatory.checkAndComplete(key), false); + TestUtils.waitUntilTrue( + () -> tryCompleteSemaphore.hasQueuedThreads(), + () -> "Not attempting to complete", + 10000, + 200); + purgatory.checkAndComplete(key); // this should not block even though lock is not free + assertFalse("Operation should not have completed", op.isCompleted()); + tryCompleteSemaphore.release(); + future.get(10, TimeUnit.SECONDS); + assertTrue("Operation should have completed", op.isCompleted()); + } + + /** + * Test `tryComplete` with multiple threads to verify that there are no timing windows + * when completion is not performed even if the thread that makes the operation completable + * may not be able to acquire the operation lock. Since it is difficult to test all scenarios, + * this test uses random delays with a large number of threads. + */ + @Test + public void testTryCompleteWithMultipleThreads() { + ScheduledExecutorService executor = Executors.newScheduledThreadPool(20); + this.executorService = executor; + Random random = ThreadLocalRandom.current(); + int maxDelayMs = 10; + final int completionAttempts = 20; + + List ops = IntStream.range(0, 100).mapToObj(index -> { + TestDelayOperation op = new TestDelayOperation(index, completionAttempts, maxDelayMs); + purgatory.tryCompleteElseWatch(op, Lists.newArrayList(op.key)); + return op; + }).collect(Collectors.toList()); + + List> futures = IntStream.rangeClosed(1, completionAttempts) + .mapToObj(i -> + ops.stream().map( + op -> scheduleTryComplete(op, random.nextInt(maxDelayMs))) + .collect(Collectors.toList())) + .flatMap(List::stream) + .collect(Collectors.toList()); + futures.forEach(future -> { + try { + future.get(); + } catch (InterruptedException | ExecutionException e) { + // no-op + } + }); + + ops.forEach(op -> assertTrue("Operation should have completed", op.isCompleted())); + } + + Future scheduleTryComplete(TestDelayOperation op, long delayMs) { + return executorService.schedule(() -> { + if (op.completionAttemptsRemaining.decrementAndGet() == 0) { + op.completable = true; + } + purgatory.checkAndComplete(op.key); + }, delayMs, TimeUnit.MILLISECONDS); + } + + @Test + public void testDelayedOperationLock() throws Exception { + verifyDelayedOperationLock(() -> new MockDelayedOperation(100000L), false); + } + + @Test + public void testDelayedOperationLockOverride() throws Exception { + verifyDelayedOperationLock(() -> { + ReentrantLock lock = new ReentrantLock(); + return new MockDelayedOperation(100000L, Optional.of(lock), Optional.of(lock)); + }, false); + + verifyDelayedOperationLock(() -> new MockDelayedOperation( + 100000L, + Optional.empty(), + Optional.of(new ReentrantLock()) + ), true); + } + + void verifyDelayedOperationLock(Supplier mockDelayedOperation, boolean mismatchedLocks) + throws Exception { + String key = "key"; + executorService = Executors.newSingleThreadScheduledExecutor(); + + Function> createDelayedOperations = count -> + IntStream.rangeClosed(1, count).mapToObj(i -> { + MockDelayedOperation op = mockDelayedOperation.get(); + purgatory.tryCompleteElseWatch(op, Lists.newArrayList(key)); + assertFalse("Not completable", op.isCompleted()); + return op; + }).collect(Collectors.toList()); + + Function> createCompletableOperations = count -> + IntStream.rangeClosed(1, count).mapToObj(i -> { + MockDelayedOperation op = mockDelayedOperation.get(); + op.completable = true; + return op; + }).collect(Collectors.toList()); + + BiFunction, List, Void> checkAndComplete = + (completableOps, expectedComplete) -> { + completableOps.forEach(op -> op.completable = true); + int completed = purgatory.checkAndComplete(key); + assertEquals(expectedComplete.size(), completed); + expectedComplete.forEach(op -> assertTrue( + "Should have completed", + op.isCompleted() + )); + Set expectedNotComplete = completableOps.stream().collect(Collectors.toSet()); + expectedComplete.forEach(op -> expectedNotComplete.remove(op)); + expectedNotComplete.forEach(op -> assertFalse("Should not have completed", op.isCompleted())); + return null; + }; + + // If locks are free all completable operations should complete + List ops = createDelayedOperations.apply(2); + checkAndComplete.apply(ops, ops); + + // Lock held by current thread, completable operations should complete + ops = createDelayedOperations.apply(2); + final List ops2 = ops; + inLock(ops.get(1).lock, () -> { + checkAndComplete.apply(ops2, ops2); + return null; + }); + + // Lock held by another thread, should not block, only operations that can be + // locked without blocking on the current thread should complete + ops = createDelayedOperations.apply(2); + final List ops3 = ops; + runOnAnotherThread(() -> ops3.get(0).lock.lock(), true); + try { + checkAndComplete.apply(ops, Lists.newArrayList(ops.get(1))); + } finally { + runOnAnotherThread(() -> ops3.get(0).lock.unlock(), true); + checkAndComplete.apply(Lists.newArrayList(ops.get(0)), Lists.newArrayList(ops.get(0))); + } + + // Lock acquired by response callback held by another thread, should not block + // if the response lock is used as operation lock, only operations + // that can be locked without blocking on the current thread should complete + ops = createDelayedOperations.apply(2); + final List ops4 = ops; + ops.get(0).responseLockOpt.map(lock -> { + try { + runOnAnotherThread(() -> lock.lock(), true); + try { + try { + checkAndComplete.apply(ops4, Lists.newArrayList(ops4.get(1))); + assertFalse("Should have failed with mismatched locks", mismatchedLocks); + } catch (IllegalStateException e) { + assertTrue("Should not have failed with valid locks", mismatchedLocks); + } + } finally { + runOnAnotherThread(() -> lock.unlock(), true); + checkAndComplete.apply(Lists.newArrayList(ops4.get(0)), Lists.newArrayList(ops4.get(0))); + } + } catch (Exception e) { + throw new RuntimeException(e); + } + return null; + }); + + // Immediately completable operations should complete without locking + ops = createCompletableOperations.apply(2); + ops.forEach(op -> { + assertTrue("Should have completed", purgatory.tryCompleteElseWatch(op, Lists.newArrayList(key))); + assertTrue("Should have completed", op.isCompleted()); + }); + } + + private Future runOnAnotherThread(Runnable f, boolean shouldComplete) throws Exception { + Future future = executorService.submit(f); + if (shouldComplete) { + future.get(); + } else { + assertFalse("Should not have completed", future.isDone()); + } + return future; + } + +} From 1087227a291f8ee38ddaae6de38ade85cab6686f Mon Sep 17 00:00:00 2001 From: Sijie Guo Date: Sat, 3 Aug 2019 23:04:34 +0800 Subject: [PATCH 6/8] Fix checkstyle and findbugs --- pom.xml | 2 +- .../coordinator/group/DelayedHeartbeat.java | 58 ++++++ .../kop/coordinator/group/DelayedJoin.java | 57 ++++++ .../kop/coordinator/group/GroupConfig.java | 4 + .../coordinator/group/GroupCoordinator.java | 166 ++++++++++-------- .../kop/coordinator/group/GroupMetadata.java | 13 +- .../group/GroupMetadataConstants.java | 28 ++- .../group/GroupMetadataManager.java | 21 ++- .../coordinator/group/InitialDelayedJoin.java | 81 +++++++++ .../io/streamnative/kop/utils/CoreUtils.java | 1 - .../group/GroupMetadataManagerTest.java | 18 +- 11 files changed, 357 insertions(+), 92 deletions(-) create mode 100644 src/main/java/io/streamnative/kop/coordinator/group/DelayedHeartbeat.java create mode 100644 src/main/java/io/streamnative/kop/coordinator/group/DelayedJoin.java create mode 100644 src/main/java/io/streamnative/kop/coordinator/group/InitialDelayedJoin.java diff --git a/pom.xml b/pom.xml index 770c9c36a1..f8b6c660f2 100644 --- a/pom.xml +++ b/pom.xml @@ -39,7 +39,7 @@ 1.18.4 2.22.0 4.1.32.Final - 2.5.0-ef274bbb8 + 2.5.0-2cc34afc0 1.7.25 3.1.8 1.11.2 diff --git a/src/main/java/io/streamnative/kop/coordinator/group/DelayedHeartbeat.java b/src/main/java/io/streamnative/kop/coordinator/group/DelayedHeartbeat.java new file mode 100644 index 0000000000..d469e2b726 --- /dev/null +++ b/src/main/java/io/streamnative/kop/coordinator/group/DelayedHeartbeat.java @@ -0,0 +1,58 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.coordinator.group; + +import io.streamnative.kop.utils.delayed.DelayedOperation; +import java.util.Optional; + +/** + * Delayed heartbeat operations that are added to the purgatory for session timeout checking. + * Heartbeats are paused during rebalance. + */ +class DelayedHeartbeat extends DelayedOperation { + + private final GroupCoordinator coordinator; + private final GroupMetadata group; + private final MemberMetadata member; + private long heartbeatDeadline; + + DelayedHeartbeat(GroupCoordinator coordinator, + GroupMetadata group, + MemberMetadata member, + long heartbeatDeadline, + long sessionTimeout) { + super(sessionTimeout, Optional.of(group.lock())); + + this.coordinator = coordinator; + this.group = group; + this.member = member; + this.heartbeatDeadline = heartbeatDeadline; + } + + @Override + public void onExpiration() { + coordinator.onExpireHeartbeat(group, member, heartbeatDeadline); + } + + @Override + public void onComplete() { + coordinator.onCompleteHeartbeat(); + } + + @Override + public boolean tryComplete() { + return coordinator.tryCompleteHeartbeat(group, member, heartbeatDeadline, () -> forceComplete()); + } + +} diff --git a/src/main/java/io/streamnative/kop/coordinator/group/DelayedJoin.java b/src/main/java/io/streamnative/kop/coordinator/group/DelayedJoin.java new file mode 100644 index 0000000000..015698770c --- /dev/null +++ b/src/main/java/io/streamnative/kop/coordinator/group/DelayedJoin.java @@ -0,0 +1,57 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.coordinator.group; + +import io.streamnative.kop.utils.delayed.DelayedOperation; +import java.util.Optional; + +/** + * Delayed rebalance operations that are added to the purgatory when group is preparing for rebalance. + * + *

Whenever a join-group request is received, check if all known group members have requested + * to re-join the group; if yes, complete this operation to proceed rebalance. + * + *

When the operation has expired, any known members that have not requested to re-join + * the group are marked as failed, and complete this operation to proceed rebalance with + * the rest of the group. + */ +class DelayedJoin extends DelayedOperation { + + final GroupCoordinator coordinator; + final GroupMetadata group; + + protected DelayedJoin(GroupCoordinator coordinator, + GroupMetadata group, + long rebalanceTimeout) { + super(rebalanceTimeout, Optional.of(group.lock())); + this.coordinator = coordinator; + this.group = group; + } + + @Override + public void onExpiration() { + coordinator.onExpireJoin(); + } + + @Override + public void onComplete() { + coordinator.onCompleteJoin(group); + } + + @Override + public boolean tryComplete() { + return coordinator.tryCompleteJoin(group, () -> forceComplete()); + } + +} diff --git a/src/main/java/io/streamnative/kop/coordinator/group/GroupConfig.java b/src/main/java/io/streamnative/kop/coordinator/group/GroupConfig.java index d609de0cef..91a185fb67 100644 --- a/src/main/java/io/streamnative/kop/coordinator/group/GroupConfig.java +++ b/src/main/java/io/streamnative/kop/coordinator/group/GroupConfig.java @@ -14,11 +14,15 @@ package io.streamnative.kop.coordinator.group; import lombok.Data; +import lombok.Getter; +import lombok.experimental.Accessors; /** * Group configuration. */ @Data +@Accessors(fluent = true) +@Getter public class GroupConfig { private final int groupMinSessionTimeoutMs; diff --git a/src/main/java/io/streamnative/kop/coordinator/group/GroupCoordinator.java b/src/main/java/io/streamnative/kop/coordinator/group/GroupCoordinator.java index adc9f65fc5..e31104af66 100644 --- a/src/main/java/io/streamnative/kop/coordinator/group/GroupCoordinator.java +++ b/src/main/java/io/streamnative/kop/coordinator/group/GroupCoordinator.java @@ -20,14 +20,16 @@ import static io.streamnative.kop.coordinator.group.GroupState.PreparingRebalance; import static io.streamnative.kop.coordinator.group.GroupState.Stable; +import com.google.common.collect.Lists; import com.google.common.collect.Sets; import io.streamnative.kop.coordinator.group.GroupMetadata.GroupOverview; import io.streamnative.kop.coordinator.group.GroupMetadata.GroupSummary; -import io.streamnative.kop.utils.CoreUtils; +import io.streamnative.kop.utils.delayed.DelayedOperationKey.GroupKey; +import io.streamnative.kop.utils.delayed.DelayedOperationKey.MemberKey; +import io.streamnative.kop.utils.delayed.DelayedOperationPurgatory; import java.util.ArrayList; import java.util.Collections; import java.util.HashMap; -import java.util.HashSet; import java.util.List; import java.util.Map; import java.util.Objects; @@ -39,7 +41,6 @@ import java.util.function.Supplier; import java.util.stream.Collectors; import lombok.extern.slf4j.Slf4j; -import lombok.val; import org.apache.kafka.common.internals.Topic; import org.apache.kafka.common.protocol.ApiKeys; import org.apache.kafka.common.protocol.Errors; @@ -47,7 +48,6 @@ import org.apache.kafka.common.utils.Time; import org.apache.pulsar.common.schema.KeyValue; import org.apache.pulsar.common.util.FutureUtil; -import scala.Array; /** * Group coordinator. @@ -90,8 +90,23 @@ private static boolean isValidGroupId(String groupId, private final AtomicBoolean isActive = new AtomicBoolean(false); private final GroupConfig groupConfig; private final GroupMetadataManager groupManager; + private final DelayedOperationPurgatory heartbeatPurgatory; + private final DelayedOperationPurgatory joinPurgatory; private final Time time; + public GroupCoordinator( + GroupConfig groupConfig, + GroupMetadataManager groupManager, + DelayedOperationPurgatory heartbeatPurgatory, + DelayedOperationPurgatory joinPurgatory, + Time time) { + this.groupConfig = groupConfig; + this.groupManager = groupManager; + this.heartbeatPurgatory = heartbeatPurgatory; + this.joinPurgatory = joinPurgatory; + this.time = time; + } + public CompletableFuture handleJoinGroup( String groupId, String memberId, @@ -108,8 +123,8 @@ public CompletableFuture handleJoinGroup( joinError(memberId, errors.get())); } - if (sessionTimeoutMs < groupConfig.getGroupMinSessionTimeoutMs() - || sessionTimeoutMs > groupConfig.getGroupMaxSessionTimeoutMs()) { + if (sessionTimeoutMs < groupConfig.groupMinSessionTimeoutMs() + || sessionTimeoutMs > groupConfig.groupMaxSessionTimeoutMs()) { return CompletableFuture.completedFuture( joinError(memberId, Errors.INVALID_SESSION_TIMEOUT)); } else { @@ -123,7 +138,7 @@ public CompletableFuture handleJoinGroup( protocolType, protocols )).orElseGet(() -> { - if (memberId != JoinGroupRequest.UNKNOWN_MEMBER_ID) { + if (!JoinGroupRequest.UNKNOWN_MEMBER_ID.equals(memberId)) { return CompletableFuture.completedFuture( joinError(memberId, Errors.UNKNOWN_MEMBER_ID)); } else { @@ -179,7 +194,7 @@ private CompletableFuture unsafeJoinGroup( ) { if (!group.is(Empty) && ( !group.protocolType().isPresent() - || group.protocolType().get() != protocolType + || !Objects.equals(group.protocolType().get(), protocolType) || !group.supportsProtocols(protocols.keySet()))) { // if the new member does not support the group protocol, reject it return CompletableFuture.completedFuture( @@ -189,7 +204,8 @@ private CompletableFuture unsafeJoinGroup( //reject if first member with empty group protocol or protocolType is empty return CompletableFuture.completedFuture( joinError(memberId, Errors.INCONSISTENT_GROUP_PROTOCOL)); - } else if (memberId != JoinGroupRequest.UNKNOWN_MEMBER_ID && !group.has(memberId)) { + } else if (!JoinGroupRequest.UNKNOWN_MEMBER_ID.equals(memberId) + && !group.has(memberId)) { // if the member trying to register with a un-recognized id, send the response to let // it reset its member id and retry return CompletableFuture.completedFuture( @@ -206,7 +222,7 @@ private CompletableFuture unsafeJoinGroup( joinError(memberId, Errors.UNKNOWN_MEMBER_ID)); break; case PreparingRebalance: - if (memberId == JoinGroupRequest.UNKNOWN_MEMBER_ID) { + if (JoinGroupRequest.UNKNOWN_MEMBER_ID.equals(memberId)) { resultFuture = addMemberAndRebalance( rebalanceTimeoutMs, sessionTimeoutMs, @@ -222,7 +238,7 @@ private CompletableFuture unsafeJoinGroup( } break; case CompletingRebalance: - if (JoinGroupRequest.UNKNOWN_MEMBER_ID == memberId) { + if (JoinGroupRequest.UNKNOWN_MEMBER_ID.equals(memberId)) { resultFuture = addMemberAndRebalance( rebalanceTimeoutMs, sessionTimeoutMs, @@ -265,7 +281,7 @@ private CompletableFuture unsafeJoinGroup( break; case Empty: case Stable: - if (JoinGroupRequest.UNKNOWN_MEMBER_ID == memberId) { + if (JoinGroupRequest.UNKNOWN_MEMBER_ID.equals(memberId)) { // if the member id is unknown, register the member to the group resultFuture = addMemberAndRebalance( rebalanceTimeoutMs, @@ -494,8 +510,8 @@ public Map handleDeleteGroups(Set groupIds) { // TODO: /// val offsetsRemoved = groupManager.cleanupGroupMetadata(groupsEligibleForDeletion, _.removeAllOffsets()) /// groupErrors ++= groupsEligibleForDeletion.map(_.groupId -> Errors.NONE).toMap - /// info(s"The following groups were deleted: ${groupsEligibleForDeletion.map(_.groupId).mkString(", ")}. " + - /// s"A total of $offsetsRemoved offsets were removed.") + /// info(s"The following groups were deleted: ${groupsEligibleForDeletion.map(_.groupId).mkString(", ")}. " + // + s"A total of $offsetsRemoved offsets were removed.") } return groupErrors; @@ -632,7 +648,7 @@ private void propagateAssignment(GroupMetadata group, Errors error) { // This is because if any member's session expired while we were still awaiting either // the leader sync group or the storage callback, its expiration will be ignored and no // future heartbeat expectations will not be scheduled. - completeAndScheduleNextHeartbeatExpiration(group, member) + completeAndScheduleNextHeartbeatExpiration(group, member); } } } @@ -648,24 +664,30 @@ private JoinGroupResult joinError(String memberId, Errors error) { } /** - * Complete existing DelayedHeartbeats for the given member and schedule the next one + * Complete existing DelayedHeartbeats for the given member and schedule the next one. */ private void completeAndScheduleNextHeartbeatExpiration(GroupMetadata group, MemberMetadata member) { // complete current heartbeat expectation member.latestHeartbeat(time.milliseconds()); - val memberKey = MemberKey(member.groupId, member.memberId) - heartbeatPurgatory.checkAndComplete(memberKey) + MemberKey memberKey = new MemberKey(member.groupId(), member.memberId()); + heartbeatPurgatory.checkAndComplete(memberKey); // reschedule the next heartbeat expiration deadline long newHeartbeatDeadline = member.latestHeartbeat() + member.sessionTimeoutMs(); - val delayedHeartbeat = new DelayedHeartbeat(this, group, member, newHeartbeatDeadline, member.sessionTimeoutMs()); - heartbeatPurgatory.tryCompleteElseWatch(delayedHeartbeat, Seq(memberKey)) + DelayedHeartbeat delayedHeartbeat = new DelayedHeartbeat( + this, + group, + member, + newHeartbeatDeadline, + member.sessionTimeoutMs()); + heartbeatPurgatory.tryCompleteElseWatch( + delayedHeartbeat, Lists.newArrayList(memberKey)); } private void removeHeartbeatForLeavingMember(GroupMetadata group, MemberMetadata member) { member.isLeaving(true); - val memberKey = MemberKey(member.groupId, member.memberId) + MemberKey memberKey = new MemberKey(member.groupId(), member.memberId()); heartbeatPurgatory.checkAndComplete(memberKey); } @@ -721,30 +743,35 @@ private void maybePrepareRebalance(GroupMetadata group) { }); } - - private void prepareRebalance(GroupMetadata group) { // if any members are awaiting sync, cancel their request and have them rejoin - if (group.is(CompletingRebalance)) - resetAndPropagateAssignmentError(group, Errors.REBALANCE_IN_PROGRESS) + if (group.is(CompletingRebalance)) { + resetAndPropagateAssignmentError(group, Errors.REBALANCE_IN_PROGRESS); + } + + DelayedJoin delayedRebalance; - val delayedRebalance = if (group.is(Empty)) - new InitialDelayedJoin(this, + if (group.is(Empty)) { + delayedRebalance = new InitialDelayedJoin(this, joinPurgatory, group, - groupConfig.groupInitialRebalanceDelayMs, - groupConfig.groupInitialRebalanceDelayMs, - max(group.rebalanceTimeoutMs - groupConfig.groupInitialRebalanceDelayMs, 0)) - else - new DelayedJoin(this, group, group.rebalanceTimeoutMs) + groupConfig.groupInitialRebalanceDelayMs(), + groupConfig.groupInitialRebalanceDelayMs(), + Math.max(group.rebalanceTimeoutMs() - groupConfig.groupInitialRebalanceDelayMs(), 0)); + } else { + delayedRebalance = new DelayedJoin(this, group, group.rebalanceTimeoutMs()); + } - group.transitionTo(PreparingRebalance) + group.transitionTo(PreparingRebalance); - info(s"Preparing to rebalance group ${group.groupId} with old generation ${group.generationId} " + - s"(${Topic.GROUP_METADATA_TOPIC_NAME}-${partitionFor(group.groupId)})") + log.info("Preparing to rebalance group {} with old generation {} ({}-{})", + group.groupId(), + group.generationId(), + Topic.GROUP_METADATA_TOPIC_NAME, + groupManager.partitionFor(group.groupId())); - val groupKey = GroupKey(group.groupId) - joinPurgatory.tryCompleteElseWatch(delayedRebalance, Seq(groupKey)) + GroupKey groupKey = new GroupKey(group.groupId()); + joinPurgatory.tryCompleteElseWatch(delayedRebalance, Lists.newArrayList(groupKey)); } private void removeMemberAndUpdateGroup(GroupMetadata group, @@ -785,7 +812,7 @@ void onCompleteJoin(GroupMetadata group) { group.inLock(() -> { // remove any members who haven't joined the group yet group.notYetRejoinedMembers().forEach(failedMember -> { - removeHeartbeatForLeavingMember(group, failedMember) + removeHeartbeatForLeavingMember(group, failedMember); group.remove(failedMember.memberId()); // TODO: cut the socket connection to the client }); @@ -797,49 +824,46 @@ void onCompleteJoin(GroupMetadata group) { group.groupId(), group.generationId(), Topic.GROUP_METADATA_TOPIC_NAME, groupManager.partitionFor(group.groupId())); - groupManager.storeGroup(group, Collections.emptyMap(), error => { + groupManager.storeGroup(group, Collections.emptyMap()).thenAccept(error -> { if (error != Errors.NONE) { - // we failed to write the empty group metadata. If the broker fails before another rebalance, - // the previous generation written to the log will become active again (and most likely timeout). - // This should be safe since there are no active members in an empty generation, so we just warn. - warn(s"Failed to write empty metadata for group ${group.groupId}: ${error.message}") + // we failed to write the empty group metadata. If the broker fails before another + // rebalance, the previous generation written to the log will become active again + // (and most likely timeout). This should be safe since there are no active members + // in an empty generation, so we just warn. + log.warn("Failed to write empty metadata for group {}: {}", + group.groupId(), error.message()); } - }) + }); } else { - info(s"Stabilized group ${group.groupId} generation ${group.generationId} " + - s"(${Topic.GROUP_METADATA_TOPIC_NAME}-${partitionFor(group.groupId)})") + log.info("Stabilized group {} generation {} ({}-{})", + group.groupId(), group.generationId(), + Topic.GROUP_METADATA_TOPIC_NAME, + groupManager.partitionFor(group.groupId())); // trigger the awaiting join group response callback for all the members after rebalancing for (MemberMetadata member : group.allMemberMetadata()) { Objects.requireNonNull(member.awaitingJoinCallback()); - JoinGroupRequest joinResult; + Map members; if (group.isLeader(member.memberId())) { - joinResult = new JoinGroupResult( - group.currentMemberMetadata - ); + members = group.currentMemberMetadata(); } else { - joinResult = new JoinGroupResult( - Collections.emptyMap() - ); + members = Collections.emptyMap(); } - = JoinGroupResult( - members = if (group.isLeader(member.memberId)) { - group.currentMemberMetadata - } else { - Map.empty - }, - memberId = member.memberId, - generationId = group.generationId, - subProtocol = group.protocolOrNull, - leaderId = group.leaderOrNull, - error = Errors.NONE) - - member.awaitingJoinCallback(joinResult) - member.awaitingJoinCallback = null - completeAndScheduleNextHeartbeatExpiration(group, member) + JoinGroupResult joinResult = new JoinGroupResult( + members, + member.memberId(), + group.generationId(), + group.protocolOrNull(), + group.leaderOrNull(), + Errors.NONE); + + member.awaitingJoinCallback().complete(joinResult); + member.awaitingJoinCallback(null); + completeAndScheduleNextHeartbeatExpiration(group, member); } } } + return null; }); } @@ -876,9 +900,9 @@ void onCompleteHeartbeat() { private boolean shouldKeepMemberAlive(MemberMetadata member, long heartbeatDeadline) { - return member.awaitingJoinCallback() != null || - member.awaitingSyncCallback() != null || - member.latestHeartbeat() + member.sessionTimeoutMs() > heartbeatDeadline; + return member.awaitingJoinCallback() != null + || member.awaitingSyncCallback() != null + || member.latestHeartbeat() + member.sessionTimeoutMs() > heartbeatDeadline; } } diff --git a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadata.java b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadata.java index 7aad590bf3..3108f859cd 100644 --- a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadata.java +++ b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadata.java @@ -39,6 +39,7 @@ import java.util.stream.Collectors; import javax.annotation.concurrent.NotThreadSafe; import lombok.Data; +import lombok.Getter; import lombok.Setter; import lombok.experimental.Accessors; import org.apache.commons.lang3.StringUtils; @@ -147,13 +148,15 @@ static class GroupSummary { } private final String groupId; + @Getter private final ReentrantLock lock = new ReentrantLock(); private GroupState state; private Optional protocolType = Optional.empty(); - private long generationId = 0L; + private int generationId = 0; private Optional leaderId = Optional.empty(); private Optional protocol = Optional.empty(); + @Getter private boolean newMemberAdded = false; // state management @@ -188,7 +191,7 @@ public String groupId() { return groupId; } - public long generationId() { + public int generationId() { return generationId; } @@ -200,6 +203,12 @@ public List allMemberMetadata() { return members.values().stream().collect(Collectors.toList()); } + public int rebalanceTimeoutMs() { + return members.values().stream().mapToInt(member -> + member.rebalanceTimeoutMs() + ).max().getAsInt(); + } + public boolean is(GroupState groupState) { return state == groupState; } diff --git a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataConstants.java b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataConstants.java index 04b18ab34b..bb0cd875f0 100644 --- a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataConstants.java +++ b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataConstants.java @@ -1,3 +1,16 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ package io.streamnative.kop.coordinator.group; import static com.google.common.base.Preconditions.checkState; @@ -15,15 +28,12 @@ import java.util.List; import java.util.Map; import java.util.stream.Collectors; -import java.util.stream.Stream; -import lombok.val; import org.apache.kafka.common.KafkaException; import org.apache.kafka.common.protocol.types.ArrayOf; import org.apache.kafka.common.protocol.types.BoundField; import org.apache.kafka.common.protocol.types.Field; import org.apache.kafka.common.protocol.types.Schema; import org.apache.kafka.common.protocol.types.Struct; -import org.apache.kafka.common.utils.Utils; import org.apache.pulsar.common.schema.KeyValue; /** @@ -152,7 +162,7 @@ final class GroupMetadataConstants { static final Schema CURRENT_OFFSET_VALUE_SCHEMA = schemaForOffset(CURRENT_OFFSET_VALUE_SCHEMA_VERSION); static final Schema CURRENT_GROUP_VALUE_SCHEMA = schemaForGroup(CURRENT_GROUP_VALUE_SCHEMA_VERSION); - private static final Schema schemaForKey(int version) { + private static Schema schemaForKey(int version) { Schema schema = MESSAGE_TYPE_SCHEMAS.get(version); if (null == schema) { throw new KafkaException("Unknown offset schema version " + version); @@ -160,7 +170,7 @@ private static final Schema schemaForKey(int version) { return schema; } - private static final Schema schemaForOffset(int version) { + private static Schema schemaForOffset(int version) { Schema schema = OFFSET_VALUE_SCHEMAS.get(version); if (null == schema) { throw new KafkaException("Unknown offset schema version " + version); @@ -168,7 +178,7 @@ private static final Schema schemaForOffset(int version) { return schema; } - private static final Schema schemaForGroup(int version) { + private static Schema schemaForGroup(int version) { Schema schema = GROUP_VALUE_SCHEMAS.get(version); if (null == schema) { throw new KafkaException("Unknown group metadata version " + version); @@ -190,7 +200,7 @@ private static Map asMap(KeyValue ...kvs) { } /** - * Generates the key for group metadata message for given group + * Generates the key for group metadata message for given group. * * @return key bytes for group metadata message */ @@ -205,7 +215,7 @@ static byte[] groupMetadataKey(String group) { /** * Generates the payload for group metadata message from given offset and metadata - * assuming the generation id, selected protocol, leader and member assignment are all available + * assuming the generation id, selected protocol, leader and member assignment are all available. * * @param groupMetadata current group metadata * @param assignment the assignment for the rebalancing generation @@ -266,7 +276,7 @@ static byte[] groupMetadataValue(GroupMetadata groupMetadata, } /** - * Decodes the offset messages' key + * Decodes the offset messages' key. */ static BaseKey readMessageKey(ByteBuffer buffer) { short version = buffer.getShort(); diff --git a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataManager.java b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataManager.java index 22ff1cf731..0b819a9b1e 100644 --- a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataManager.java +++ b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataManager.java @@ -1,3 +1,16 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ package io.streamnative.kop.coordinator.group; import static io.streamnative.kop.coordinator.group.GroupMetadataConstants.CURRENT_GROUP_VALUE_SCHEMA_VERSION; @@ -102,7 +115,7 @@ public String toString() { private final ReentrantLock partitionLock = new ReentrantLock(); /** * partitions of consumer groups that are being loaded, its lock should - * be always called BEFORE the group lock if needed + * be always called BEFORE the group lock if needed. */ private final Set loadingPartitions = new HashSet<>(); /* partitions of consumer groups that are assigned, using the same loading partition lock */ @@ -378,9 +391,9 @@ private void loadGroup(GroupMetadata group) { } /** - * Add the partition into the owned list + * Add the partition into the owned list. * - * NOTE: this is for test only + *

NOTE: this is for test only */ private void addPartitionOwnership(int partition) { inLock(partitionLock, () -> { @@ -393,7 +406,7 @@ private void addPartitionOwnership(int partition) { * Add a partition to the loading partitions set. Return true if the partition was not * already loading. * - * Visible for testing + *

Visible for testing */ boolean addLoadingPartition(int partition) { return inLock(partitionLock, () -> loadingPartitions.add(partition)); diff --git a/src/main/java/io/streamnative/kop/coordinator/group/InitialDelayedJoin.java b/src/main/java/io/streamnative/kop/coordinator/group/InitialDelayedJoin.java new file mode 100644 index 0000000000..5e3eb239b7 --- /dev/null +++ b/src/main/java/io/streamnative/kop/coordinator/group/InitialDelayedJoin.java @@ -0,0 +1,81 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.coordinator.group; + +import com.google.common.collect.Lists; +import io.streamnative.kop.utils.delayed.DelayedOperationKey.GroupKey; +import io.streamnative.kop.utils.delayed.DelayedOperationPurgatory; + +/** + * Delayed rebalance operation that is added to the purgatory when a group is transitioning from + * Empty to PreparingRebalance + * + *

When onComplete is triggered we check if any new members have been added and if there is still time remaining + * before the rebalance timeout. If both are true we then schedule a further delay. Otherwise we complete the + * rebalance. + */ +class InitialDelayedJoin extends DelayedJoin { + + final DelayedOperationPurgatory purgatory; + final int configuredRebalanceDelay; + final int delayMs; + final int remainingMs; + + InitialDelayedJoin(GroupCoordinator coordinator, + DelayedOperationPurgatory purgatory, + GroupMetadata group, + int configuredRebalanceDelay, + int delayMs, + int remainingMs) { + super(coordinator, group, delayMs); + this.purgatory = purgatory; + this.configuredRebalanceDelay = configuredRebalanceDelay; + this.delayMs = delayMs; + this.remainingMs = remainingMs; + } + + @Override + public void onComplete() { + group.inLock(() -> { + if (group.newMemberAdded() && remainingMs != 0) { + group.newMemberAdded(false); + int delay = Math.min(configuredRebalanceDelay, remainingMs); + int remaining = Math.max( + remainingMs - delayMs, + 0 + ); + purgatory.tryCompleteElseWatch( + new InitialDelayedJoin( + coordinator, + purgatory, + group, + configuredRebalanceDelay, + delay, + remaining + ), + Lists.newArrayList(new GroupKey(group.groupId())) + ); + } else { + super.onComplete(); + } + return null; + }); + } + + @Override + public boolean tryComplete() { + return false; + } + +} diff --git a/src/main/java/io/streamnative/kop/utils/CoreUtils.java b/src/main/java/io/streamnative/kop/utils/CoreUtils.java index 9af289ee11..3987086b31 100644 --- a/src/main/java/io/streamnative/kop/utils/CoreUtils.java +++ b/src/main/java/io/streamnative/kop/utils/CoreUtils.java @@ -41,5 +41,4 @@ public static T inWriteLock(ReadWriteLock lock, Supplier supplier) { return inLock(lock.writeLock(), supplier); } - private CoreUtils() {} } diff --git a/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataManagerTest.java b/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataManagerTest.java index 296c8ebb2f..da10a977cb 100644 --- a/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataManagerTest.java +++ b/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataManagerTest.java @@ -1,10 +1,20 @@ -functionspackage io.streamnative.kop.coordinator.group; +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.coordinator.group; /** * Unit test {@link GroupMetadataManager}. */ public class GroupMetadataManagerTest { - - - } From 5c2124feb760ed1fbdb0497436ad0876a24eb096 Mon Sep 17 00:00:00 2001 From: Sijie Guo Date: Sun, 4 Aug 2019 10:11:27 +0800 Subject: [PATCH 7/8] Add GroupMetadataManagerTest --- .../group/GroupMetadataManager.java | 35 +-- .../group/GroupMetadataManagerTest.java | 277 +++++++++++++++++- 2 files changed, 282 insertions(+), 30 deletions(-) diff --git a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataManager.java b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataManager.java index 0b819a9b1e..ba151240ab 100644 --- a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataManager.java +++ b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataManager.java @@ -41,8 +41,6 @@ import org.apache.bookkeeper.common.util.MathUtils; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.protocol.Errors; -import org.apache.kafka.common.record.CompressionType; -import org.apache.kafka.common.requests.ApiVersionsResponse.ApiVersion; import org.apache.kafka.common.utils.Time; import org.apache.pulsar.client.api.MessageId; import org.apache.pulsar.client.api.Producer; @@ -106,10 +104,6 @@ public String toString() { } - private final int brokerId; - private final ApiVersion interBrokerProtocolVersion; - private final OffsetConfig config; - private final CompressionType compressionType; private final ConcurrentMap groupMetadataCache; /* lock protecting access to loading and owned partition sets */ private final ReentrantLock partitionLock = new ReentrantLock(); @@ -122,25 +116,16 @@ public String toString() { private final Set ownedPartitions = new HashSet<>(); /* shutting down flag */ private final AtomicBoolean shuttingDown = new AtomicBoolean(false); - private final String logIdent; private final int groupMetadataTopicPartitionCount; private final Producer metadataTopicProducer; private final Reader metadataTopicReader; private final Time time; - GroupMetadataManager(int brokerId, - ApiVersion interBrokerProtocolVersion, - OffsetConfig config, - int groupMetadataTopicPartitionCount, + GroupMetadataManager(int groupMetadataTopicPartitionCount, Producer metadataTopicProducer, Reader metadataTopicConsumer, Time time) { - this.brokerId = brokerId; - this.interBrokerProtocolVersion = interBrokerProtocolVersion; - this.config = config; - this.compressionType = config.offsetsTopicCompressionType(); this.groupMetadataCache = new ConcurrentHashMap<>(); - this.logIdent = String.format("[GroupMetadataManager brokerId=%d]", brokerId); this.groupMetadataTopicPartitionCount = groupMetadataTopicPartitionCount; this.metadataTopicProducer = metadataTopicProducer; this.metadataTopicReader = metadataTopicConsumer; @@ -191,15 +176,10 @@ public boolean isLoading() { public boolean groupNotExists(String groupId) { return inLock( partitionLock, - () -> { - if (isGroupLocal(groupId)) { - return true; - } else { - return getGroup(groupId) - .map(metadata -> metadata.inLock(() -> metadata.is(GroupState.Dead))) - .orElse(false); - } - } + () -> isGroupLocal(groupId) + && getGroup(groupId) + .map(group -> group.inLock(() -> group.is(GroupState.Dead))) + .orElse(true) ); } @@ -395,7 +375,7 @@ private void loadGroup(GroupMetadata group) { * *

NOTE: this is for test only */ - private void addPartitionOwnership(int partition) { + void addPartitionOwnership(int partition) { inLock(partitionLock, () -> { ownedPartitions.add(partition); return null; @@ -412,7 +392,4 @@ boolean addLoadingPartition(int partition) { return inLock(partitionLock, () -> loadingPartitions.add(partition)); } - - - } diff --git a/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataManagerTest.java b/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataManagerTest.java index da10a977cb..3b2e4b3911 100644 --- a/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataManagerTest.java +++ b/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataManagerTest.java @@ -13,8 +13,283 @@ */ package io.streamnative.kop.coordinator.group; +import static org.junit.Assert.assertEquals; +import static org.junit.Assert.assertFalse; +import static org.junit.Assert.assertTrue; + +import com.google.common.collect.Sets; +import io.streamnative.kop.MockKafkaServiceBaseTest; +import io.streamnative.kop.coordinator.group.GroupMetadataManager.BaseKey; +import io.streamnative.kop.coordinator.group.GroupMetadataManager.GroupMetadataKey; +import io.streamnative.kop.utils.MockTime; +import java.nio.ByteBuffer; +import java.util.Collections; +import java.util.HashMap; +import java.util.Map; +import java.util.Optional; +import java.util.concurrent.CompletableFuture; +import lombok.Cleanup; +import lombok.extern.slf4j.Slf4j; +import org.apache.kafka.common.TopicPartition; +import org.apache.kafka.common.internals.Topic; +import org.apache.kafka.common.protocol.Errors; +import org.apache.pulsar.client.api.Consumer; +import org.apache.pulsar.client.api.Message; +import org.apache.pulsar.client.api.MessageId; +import org.apache.pulsar.client.api.Producer; +import org.apache.pulsar.client.api.Reader; +import org.apache.pulsar.client.api.SubscriptionInitialPosition; +import org.apache.pulsar.common.policies.data.ClusterData; +import org.apache.pulsar.common.policies.data.RetentionPolicies; +import org.apache.pulsar.common.policies.data.TenantInfo; +import org.junit.After; +import org.junit.Before; +import org.junit.Test; + /** * Unit test {@link GroupMetadataManager}. */ -public class GroupMetadataManagerTest { +@Slf4j +public class GroupMetadataManagerTest extends MockKafkaServiceBaseTest { + + private static final String groupId = "foo"; + private static final int groupPartitionId = 0; + private static final TopicPartition groupTopicPartition = + new TopicPartition(Topic.GROUP_METADATA_TOPIC_NAME, groupPartitionId); + private static final String protocolType = "protocolType"; + private static final int rebalanceTimeout = 60000; + private static final int sessionTimeout = 10000; + + MockTime time = null; + GroupMetadataManager groupMetadataManager = null; + Producer producer = null; + Reader consumer = null; + + @Before + @Override + public void setup() throws Exception { + super.internalSetup(); + log.info("Admin : {}", admin); + admin.clusters().createCluster("test", + new ClusterData("http://127.0.0.1:" + brokerWebservicePort)); + + admin.tenants().createTenant("public", + new TenantInfo(Sets.newHashSet("appid1", "appid2"), Sets.newHashSet("test"))); + admin.namespaces().createNamespace("public/default"); + admin.namespaces().setNamespaceReplicationClusters("public/default", Sets.newHashSet("test")); + admin.namespaces().setRetention("public/default", + new RetentionPolicies(20, 100)); + + time = new MockTime(); + groupMetadataManager = new GroupMetadataManager( + 1, + producer, + consumer, + time + ); + } + + @After + @Override + public void cleanup() throws Exception { + super.internalCleanup(); + } + + @Test + public void testGroupNotExits() { + // group is not owned + assertFalse(groupMetadataManager.groupNotExists(groupId)); + + groupMetadataManager.addPartitionOwnership(groupPartitionId); + // group is owned but does not exist yet + assertTrue(groupMetadataManager.groupNotExists(groupId)); + + GroupMetadata group = new GroupMetadata(groupId, GroupState.Empty); + groupMetadataManager.addGroup(group); + + // group is owned but not Dead + assertFalse(groupMetadataManager.groupNotExists(groupId)); + + group.transitionTo(GroupState.Dead); + // group is owned and Dead + assertTrue(groupMetadataManager.groupNotExists(groupId)); + } + + @Test + public void testAddGroup() { + GroupMetadata group = new GroupMetadata("foo", GroupState.Empty); + assertEquals(group, groupMetadataManager.addGroup(group)); + assertEquals(group, groupMetadataManager.addGroup( + new GroupMetadata("foo", GroupState.Empty) + )); + } + + /** + * A group metadata manager test runner. + */ + @FunctionalInterface + public interface GroupMetadataManagerTester { + + void test(GroupMetadataManager groupMetadataManager, + Consumer consumer) throws Exception; + + } + + void runGroupMetadataManagerTester(final String topicName, + GroupMetadataManagerTester tester) throws Exception { + @Cleanup + Producer producer = pulsarClient.newProducer() + .topic(topicName) + .create(); + @Cleanup + Consumer consumer = pulsarClient.newConsumer() + .topic(topicName) + .subscriptionName("test-sub") + .subscriptionInitialPosition(SubscriptionInitialPosition.Earliest) + .subscribe(); + @Cleanup + Reader reader = pulsarClient.newReader() + .topic(topicName) + .startMessageId(MessageId.earliest) + .create(); + groupMetadataManager = new GroupMetadataManager( + 1, + producer, + reader, + time + ); + tester.test(groupMetadataManager, consumer); + } + + @Test + public void testStoreEmptyGroup() throws Exception { + final String topicName = "test-store-empty-group"; + + runGroupMetadataManagerTester(topicName, (groupMetadataManager, consumer) -> { + int generation = 27; + String protocolType = "consumer"; + GroupMetadata group = GroupMetadata.loadGroup( + groupId, + GroupState.Empty, + generation, + protocolType, + null, + null, + Collections.emptyList() + ); + groupMetadataManager.addGroup(group); + + Errors errors = groupMetadataManager.storeGroup(group, Collections.emptyMap()).get(); + assertEquals(Errors.NONE, errors); + + Message message = consumer.receive(); + assertTrue(message.getEventTime() > 0L); + assertTrue(message.hasKey()); + byte[] key = message.getKeyBytes(); + byte[] value = message.getValue(); + + BaseKey bk = GroupMetadataConstants.readMessageKey(ByteBuffer.wrap(key)); + assertTrue(bk instanceof GroupMetadataKey); + GroupMetadataKey gmk = (GroupMetadataKey) bk; + assertEquals(groupId, gmk.key()); + + GroupMetadata gm = GroupMetadataConstants.readGroupMessageValue( + groupId, ByteBuffer.wrap(value) + ); + assertTrue(gm.is(GroupState.Empty)); + assertEquals(generation, gm.generationId()); + assertEquals(Optional.of(protocolType), gm.protocolType()); + }); + } + + @Test + public void testStoreEmptySimpleGroup() throws Exception { + final String topicName = "test-store-empty-simple-group"; + + runGroupMetadataManagerTester(topicName, (groupMetadataManager, consumer) -> { + + GroupMetadata group = new GroupMetadata(groupId, GroupState.Empty); + groupMetadataManager.addGroup(group); + + Errors errors = groupMetadataManager.storeGroup(group, Collections.emptyMap()).get(); + assertEquals(Errors.NONE, errors); + + Message message = consumer.receive(); + assertTrue(message.getEventTime() > 0L); + assertTrue(message.hasKey()); + byte[] key = message.getKeyBytes(); + byte[] value = message.getValue(); + + BaseKey bk = GroupMetadataConstants.readMessageKey(ByteBuffer.wrap(key)); + assertTrue(bk instanceof GroupMetadataKey); + GroupMetadataKey gmk = (GroupMetadataKey) bk; + assertEquals(groupId, gmk.key()); + + GroupMetadata gm = GroupMetadataConstants.readGroupMessageValue( + groupId, ByteBuffer.wrap(value) + ); + assertTrue(gm.is(GroupState.Empty)); + assertEquals(0, gm.generationId()); + assertEquals(Optional.empty(), gm.protocolType()); + + }); + } + + @Test + public void testStoreNoneEmptyGroup() throws Exception { + final String topicName = "test-store-non-empty-group"; + + runGroupMetadataManagerTester(topicName, (groupMetadataManager, consumer) -> { + String memberId = "memberId"; + String clientId = "clientId"; + String clientHost = "localhost"; + + GroupMetadata group = new GroupMetadata(groupId, GroupState.Empty); + groupMetadataManager.addGroup(group); + + Map protocols = new HashMap<>(); + protocols.put("protocol", new byte[0]); + MemberMetadata member = new MemberMetadata( + memberId, + groupId, + clientId, + clientHost, + rebalanceTimeout, + sessionTimeout, + protocolType, + protocols + ); + CompletableFuture joinFuture = new CompletableFuture<>(); + member.awaitingJoinCallback(joinFuture); + group.add(member); + group.transitionTo(GroupState.PreparingRebalance); + group.initNextGeneration(); + + Map assignments = new HashMap<>(); + assignments.put(memberId, new byte[0]); + Errors errors = groupMetadataManager.storeGroup(group, assignments).get(); + assertEquals(Errors.NONE, errors); + + Message message = consumer.receive(); + assertTrue(message.getEventTime() > 0L); + assertTrue(message.hasKey()); + byte[] key = message.getKeyBytes(); + byte[] value = message.getValue(); + + BaseKey bk = GroupMetadataConstants.readMessageKey(ByteBuffer.wrap(key)); + assertTrue(bk instanceof GroupMetadataKey); + GroupMetadataKey gmk = (GroupMetadataKey) bk; + assertEquals(groupId, gmk.key()); + + GroupMetadata gm = GroupMetadataConstants.readGroupMessageValue( + groupId, ByteBuffer.wrap(value) + ); + assertEquals(GroupState.Stable, gm.currentState()); + assertEquals(1, gm.generationId()); + assertEquals(Optional.of(protocolType), gm.protocolType()); + assertEquals("protocol", gm.protocolOrNull()); + assertTrue(gm.has(memberId)); + }); + } + } From fff811027191a239a84915fc2243079e42d95478 Mon Sep 17 00:00:00 2001 From: Sijie Guo Date: Sun, 4 Aug 2019 12:24:35 +0800 Subject: [PATCH 8/8] Add offset metadata to group metadata *Motivation* Add the support for offsets --- .../kop/coordinator/group/GroupMetadata.java | 245 ++++++++++++++++++ .../group/GroupMetadataManager.java | 62 +++++ .../kop/coordinator/group/OffsetConfig.java | 8 +- .../kop/offset/OffsetAndMetadata.java | 103 ++++++++ .../kop/offset/OffsetMetadata.java | 52 ++++ .../kop/offset/OffsetMetadataAndError.java | 89 +++++++ .../streamnative/kop/offset/package-info.java | 17 ++ .../group/GroupMetadataManagerTest.java | 3 + .../coordinator/group/GroupMetadataTest.java | 216 ++++++++++++++- 9 files changed, 793 insertions(+), 2 deletions(-) create mode 100644 src/main/java/io/streamnative/kop/offset/OffsetAndMetadata.java create mode 100644 src/main/java/io/streamnative/kop/offset/OffsetMetadata.java create mode 100644 src/main/java/io/streamnative/kop/offset/OffsetMetadataAndError.java create mode 100644 src/main/java/io/streamnative/kop/offset/package-info.java diff --git a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadata.java b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadata.java index 3108f859cd..14e848b062 100644 --- a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadata.java +++ b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadata.java @@ -23,6 +23,7 @@ import com.google.common.base.Supplier; import com.google.common.collect.Sets; import io.streamnative.kop.coordinator.group.MemberMetadata.MemberSummary; +import io.streamnative.kop.offset.OffsetAndMetadata; import io.streamnative.kop.utils.CoreUtils; import java.util.Collections; import java.util.Comparator; @@ -37,12 +38,16 @@ import java.util.UUID; import java.util.concurrent.locks.ReentrantLock; import java.util.stream.Collectors; +import java.util.stream.Stream; import javax.annotation.concurrent.NotThreadSafe; import lombok.Data; import lombok.Getter; import lombok.Setter; import lombok.experimental.Accessors; +import lombok.extern.slf4j.Slf4j; import org.apache.commons.lang3.StringUtils; +import org.apache.kafka.common.TopicPartition; +import org.apache.pulsar.common.schema.KeyValue; /** * Group contains the following metadata: @@ -60,6 +65,7 @@ @NotThreadSafe @Setter @Accessors(fluent = true) +@Slf4j class GroupMetadata { private static final Map> validPreviousStates = new HashMap<>(); @@ -147,6 +153,22 @@ static class GroupSummary { private final List members; } + /** + * We cache offset commits along with their commit record offset. This enables us to ensure that the latest offset + * commit is always materialized when we have a mix of transactional and regular offset commits. Without preserving + * information of the commit record offset, compaction of the offsets topic it self may result in the wrong offset + * commit being materialized. + */ + @Data + static class CommitRecordMetadataAndOffset { + private final Optional appendedBatchOffset; + private final OffsetAndMetadata offsetAndMetadata; + + public boolean olderThan(CommitRecordMetadataAndOffset that) { + return appendedBatchOffset.get() < that.appendedBatchOffset.get(); + } + } + private final String groupId; @Getter private final ReentrantLock lock = new ReentrantLock(); @@ -161,6 +183,12 @@ static class GroupSummary { // state management private final Map members = new HashMap<>(); + private final Map offsets = new HashMap<>(); + private final Map pendingOffsetCommits = new HashMap<>(); + private final Map> pendingTransactionalOffsetCommits = + new HashMap<>(); + private boolean receivedTransactionalOffsetCommits = false; + private boolean receivedConsumerOffsetCommits = false; GroupMetadata(String groupId, GroupState initialState) { this.groupId = groupId; @@ -389,6 +417,223 @@ public GroupOverview overview() { ); } + public void initializeOffsets(Map offsets, + Map> pendingTxnOffsets) { + this.offsets.putAll(offsets); + this.pendingTransactionalOffsetCommits.putAll(pendingTxnOffsets); + } + + public void onOffsetCommitAppend(TopicPartition topicPartition, + CommitRecordMetadataAndOffset offsetWithCommitRecordMetadata) { + if (pendingOffsetCommits.containsKey(topicPartition)) { + if (!offsetWithCommitRecordMetadata.appendedBatchOffset.isPresent()) { + throw new IllegalStateException("Cannot complete offset commit write without providing the metadata" + + " of the record in the log."); + } + if (!offsets.containsKey(topicPartition) + || offsets.get(topicPartition).olderThan(offsetWithCommitRecordMetadata)) { + offsets.put(topicPartition, offsetWithCommitRecordMetadata); + } + } + + OffsetAndMetadata stagedOffset = pendingOffsetCommits.get(topicPartition); + if (null != stagedOffset && offsetWithCommitRecordMetadata.offsetAndMetadata == stagedOffset) { + pendingOffsetCommits.remove(topicPartition); + } else { + // The pendingOffsetCommits for this partition could be empty if the topic was deleted, in which case + // its entries would be removed from the cache by the `removeOffsets` method. + } + } + + public void failPendingOffsetWrite(TopicPartition topicPartition, + OffsetAndMetadata offset) { + OffsetAndMetadata pendingOffset = pendingOffsetCommits.get(topicPartition); + if (pendingOffset != null && offset == pendingOffset) { + pendingOffsetCommits.remove(topicPartition); + } + } + + public void prepareOffsetCommit(Map offsets) { + receivedConsumerOffsetCommits = true; + pendingOffsetCommits.putAll(offsets); + } + + public void prepareTxnOffsetCommit(long producerId, + Map offsets) { + if (log.isTraceEnabled()) { + log.trace("TxnOffsetCommit for producer {} and group {} with offsets {} is pending", + producerId, groupId, offsets); + } + receivedTransactionalOffsetCommits = true; + Map producerOffsets = + pendingTransactionalOffsetCommits.computeIfAbsent(producerId, pid -> new HashMap<>()); + offsets.forEach((tp, offsetsAndMetadata) -> producerOffsets.put(tp, new CommitRecordMetadataAndOffset( + Optional.empty(), + offsetsAndMetadata + ))); + } + + public boolean hasReceivedConsistentOffsetCommits() { + return !receivedConsumerOffsetCommits || !receivedTransactionalOffsetCommits; + } + + /** + * Remove a pending transactional offset commit if the actual offset commit record was not written to the log. + * We will return an error and the client will retry the request, potentially to a different coordinator. + */ + public void failPendingTxnOffsetCommit(long producerId, + TopicPartition topicPartition) { + Map pendingOffsets = + pendingTransactionalOffsetCommits.get(producerId); + if (null != pendingOffsets) { + CommitRecordMetadataAndOffset pendingOffsetCommit = pendingOffsets.remove(topicPartition); + if (log.isTraceEnabled()) { + log.trace("TxnOffsetCommit for producer {} and group {} with offsets {} failed to be appended" + + " to the log", + producerId, groupId, pendingOffsetCommit); + } + if (pendingOffsets.isEmpty()) { + pendingTransactionalOffsetCommits.remove(producerId); + } + } else { + // We may hit this case if the partition in question has emigrated already. + } + } + + public void onTxnOffsetCommitAppend(long producerId, + TopicPartition topicPartition, + CommitRecordMetadataAndOffset commitRecordMetadataAndOffset) { + Map pendingOffsets = + pendingTransactionalOffsetCommits.get(producerId); + if (null != pendingOffsets) { + if (pendingOffsets.containsKey(topicPartition) + && pendingOffsets.get(topicPartition).offsetAndMetadata() + == commitRecordMetadataAndOffset.offsetAndMetadata) { + pendingOffsets.put(topicPartition, commitRecordMetadataAndOffset); + } + } else { + // We may hit this case if the partition in question has emigrated. + } + } + + /** + * Complete a pending transactional offset commit. This is called after a commit or abort marker is fully written + * to the log. + */ + public void completePendingTxnOffsetCommit(long producerId, + boolean isCommit) { + Map pendingOffsets = + pendingTransactionalOffsetCommits.remove(producerId); + if (isCommit) { + if (null != pendingOffsets) { + pendingOffsets.entrySet().forEach(e -> { + TopicPartition topicPartition = e.getKey(); + CommitRecordMetadataAndOffset commitRecordMetadataAndOffset = e.getValue(); + if (!commitRecordMetadataAndOffset.appendedBatchOffset.isPresent()) { + throw new IllegalStateException(String.format("Trying to complete a transactional offset" + + " commit for producerId %s and groupId %s even though the offset commit record" + + " itself hasn't been appended to the log.", producerId, groupId)); + } + + CommitRecordMetadataAndOffset currentOffsetOpt = offsets.get(topicPartition); + if (currentOffsetOpt == null || currentOffsetOpt.olderThan(commitRecordMetadataAndOffset)) { + if (log.isTraceEnabled()) { + log.trace("TxnOffsetCommit for producer {} and group {} with offset {} " + + "committed and loaded into the cache.", + producerId, groupId, commitRecordMetadataAndOffset); + } + offsets.put(topicPartition, commitRecordMetadataAndOffset); + } else { + if (log.isTraceEnabled()) { + log.trace("TxnOffsetCommit for producer {} and group {} with offset {} " + + "committed, but not loaded since its offset is older than current offset" + + " {}.", + producerId, groupId, commitRecordMetadataAndOffset, currentOffsetOpt); + } + } + }); + } + } else { + if (log.isTraceEnabled()) { + log.trace("TxnOffsetCommit for producer {} and group {} with offsets {} aborted", + producerId, groupId, pendingOffsets); + } + } + } + + public Set activeProducers() { + return pendingTransactionalOffsetCommits.keySet(); + } + + public boolean hasPendingOffsetCommitsFromProducer(long producerId) { + return pendingTransactionalOffsetCommits.containsKey(producerId); + } + + public Map removeAllOffsets() { + return removeOffsets(offsets.keySet().stream()); + } + + public Map removeOffsets(Stream topicPartitions) { + return topicPartitions.map(topicPartition -> { + pendingOffsetCommits.remove(topicPartition); + pendingTransactionalOffsetCommits.forEach((pid, pendingOffsets) -> { + pendingOffsets.remove(topicPartition); + }); + CommitRecordMetadataAndOffset removedOffset = offsets.remove(topicPartition); + return new KeyValue<>( + topicPartition, + removedOffset.offsetAndMetadata() + ); + }).collect(Collectors.toMap( + e -> e.getKey(), + e -> e.getValue() + )); + } + + public Map removeExpiredOffsets(long startMs) { + Map expiredOffsets = offsets.entrySet().stream() + .filter(e -> + e.getValue().offsetAndMetadata().expireTimestamp() < startMs + && !pendingOffsetCommits.containsKey(e.getKey())) + .map(e -> new KeyValue<>( + e.getKey(), + e.getValue().offsetAndMetadata() + )) + .collect(Collectors.toMap( + kv -> kv.getKey(), + kv -> kv.getValue() + )); + + expiredOffsets.keySet().forEach(tp -> offsets.remove(tp)); + return expiredOffsets; + } + + public Map allOffsets() { + return offsets.entrySet().stream().collect(Collectors.toMap( + e -> e.getKey(), + e -> e.getValue().offsetAndMetadata() + )); + } + + public Optional offset(TopicPartition topicPartition) { + return Optional.ofNullable(offsets.get(topicPartition)).map(e -> e.offsetAndMetadata); + } + + // visible for testing + Optional offsetWithRecordMetadata(TopicPartition topicPartition) { + return Optional.ofNullable(offsets.get(topicPartition)); + } + + public int numOffsets() { + return offsets.size(); + } + + public boolean hasOffsets() { + return !offsets.isEmpty() + || !pendingOffsetCommits.isEmpty() + || !pendingTransactionalOffsetCommits.isEmpty(); + } + @Override public String toString() { ToStringHelper helper = MoreObjects.toStringHelper("GroupMetadata") diff --git a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataManager.java b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataManager.java index ba151240ab..2eedbda527 100644 --- a/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataManager.java +++ b/src/main/java/io/streamnative/kop/coordinator/group/GroupMetadataManager.java @@ -22,7 +22,9 @@ import static java.nio.charset.StandardCharsets.UTF_8; import static org.apache.kafka.common.internals.Topic.GROUP_METADATA_TOPIC_NAME; +import io.streamnative.kop.offset.OffsetAndMetadata; import java.nio.ByteBuffer; +import java.util.Collections; import java.util.HashMap; import java.util.HashSet; import java.util.Map; @@ -41,6 +43,7 @@ import org.apache.bookkeeper.common.util.MathUtils; import org.apache.kafka.common.TopicPartition; import org.apache.kafka.common.protocol.Errors; +import org.apache.kafka.common.record.RecordBatch; import org.apache.kafka.common.utils.Time; import org.apache.pulsar.client.api.MessageId; import org.apache.pulsar.client.api.Producer; @@ -104,6 +107,7 @@ public String toString() { } + private final OffsetConfig config; private final ConcurrentMap groupMetadataCache; /* lock protecting access to loading and owned partition sets */ private final ReentrantLock partitionLock = new ReentrantLock(); @@ -117,14 +121,25 @@ public String toString() { /* shutting down flag */ private final AtomicBoolean shuttingDown = new AtomicBoolean(false); private final int groupMetadataTopicPartitionCount; + + /** + * The groups with open transactional offsets commits per producer. We need this because when the commit or abort + * marker comes in for a transaction, it is for a particular partition on the offsets topic and a particular + * producerId. We use this structure to quickly find the groups which need to be updated by the commit/abort + * marker. + */ + private final Map> openGroupsForProducer = new HashMap<>(); + private final Producer metadataTopicProducer; private final Reader metadataTopicReader; private final Time time; GroupMetadataManager(int groupMetadataTopicPartitionCount, + OffsetConfig config, Producer metadataTopicProducer, Reader metadataTopicConsumer, Time time) { + this.config = config; this.groupMetadataCache = new ConcurrentHashMap<>(); this.groupMetadataTopicPartitionCount = groupMetadataTopicPartitionCount; this.metadataTopicProducer = metadataTopicProducer; @@ -183,6 +198,13 @@ && getGroup(groupId) ); } + boolean isGroupOpenForProducer(long producerId, + String groupId) { + return openGroupsForProducer.getOrDefault( + producerId, Collections.emptySet() + ).contains(groupId); + } + public Optional getGroup(String groupId) { return Optional.ofNullable(groupMetadataCache.getOrDefault(groupId, null)); } @@ -212,6 +234,46 @@ public CompletableFuture storeGroup(GroupMetadata group, .exceptionally(cause -> Errors.COORDINATOR_NOT_AVAILABLE); } + public CompletableFuture> storeOffsets( + GroupMetadata group, + String consumerId, + Map offsetMetadata + ) { + return storeOffsets( + group, + consumerId, + offsetMetadata, + RecordBatch.NO_PRODUCER_ID, + RecordBatch.NO_PRODUCER_EPOCH + ); + } + + public CompletableFuture> storeOffsets( + GroupMetadata group, + String consumerId, + Map offsetMetadata, + long producerId, + short producerEpoch + ) { + // first filter out partitions with offset metadata size exceeding limit + // Map filteredOffsetMetadata = + // offsetMetadata.entrySet().stream() + // .filter(entry -> validateOffsetMetadataLength(entry.getValue().metadata())) + // .collect(Collectors.toMap( + // e -> e.getKey(), + // e -> e.getValue() + // )); + + throw new UnsupportedOperationException(); + } + + /* + * Check if the offset metadata length is valid + */ + private boolean validateOffsetMetadataLength(String metadata) { + return metadata == null || metadata.length() <= config.maxMetadataSize(); + } + public CompletableFuture scheduleLoadGroupAndOffsets(int offsetsPartition, Consumer onGroupLoaded) { TopicPartition topicPartition = new TopicPartition( diff --git a/src/main/java/io/streamnative/kop/coordinator/group/OffsetConfig.java b/src/main/java/io/streamnative/kop/coordinator/group/OffsetConfig.java index cbf172480f..aaf4ba4c94 100644 --- a/src/main/java/io/streamnative/kop/coordinator/group/OffsetConfig.java +++ b/src/main/java/io/streamnative/kop/coordinator/group/OffsetConfig.java @@ -14,6 +14,7 @@ package io.streamnative.kop.coordinator.group; import lombok.Builder; +import lombok.Builder.Default; import lombok.Data; import lombok.experimental.Accessors; import org.apache.kafka.common.record.CompressionType; @@ -26,6 +27,11 @@ @Accessors(fluent = true) public class OffsetConfig { - private final CompressionType offsetsTopicCompressionType; + private static final int DefaultMaxMetadataSize = 4096; + + @Default + private int maxMetadataSize = DefaultMaxMetadataSize; + @Default + private CompressionType offsetsTopicCompressionType = CompressionType.NONE; } diff --git a/src/main/java/io/streamnative/kop/offset/OffsetAndMetadata.java b/src/main/java/io/streamnative/kop/offset/OffsetAndMetadata.java new file mode 100644 index 0000000000..da7a5857f5 --- /dev/null +++ b/src/main/java/io/streamnative/kop/offset/OffsetAndMetadata.java @@ -0,0 +1,103 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.offset; + +import static org.apache.kafka.common.requests.OffsetCommitRequest.DEFAULT_TIMESTAMP; + +import lombok.AccessLevel; +import lombok.Data; +import lombok.RequiredArgsConstructor; +import lombok.experimental.Accessors; + +/** + * Offset and metadata. + */ +@Data +@Accessors(fluent = true) +@RequiredArgsConstructor(access = AccessLevel.PRIVATE) +public class OffsetAndMetadata { + + public static OffsetAndMetadata apply( + long offset, + String metadata, + long commitTimestamp, + long expireTimestamp + ) { + return new OffsetAndMetadata( + new OffsetMetadata(offset, metadata), + commitTimestamp, + expireTimestamp + ); + } + + public static OffsetAndMetadata apply( + long offset, + String metadata, + long timestamp + ) { + return new OffsetAndMetadata( + new OffsetMetadata(offset, metadata), + timestamp, + timestamp + ); + } + + public static OffsetAndMetadata apply( + long offset, + String metadata + ) { + return new OffsetAndMetadata( + new OffsetMetadata(offset, metadata) + ); + } + + public static OffsetAndMetadata apply( + long offset + ) { + return new OffsetAndMetadata( + new OffsetMetadata(offset, OffsetMetadata.NO_METADATA) + ); + } + + private final OffsetMetadata offsetMetadata; + private final long commitTimestamp; + private final long expireTimestamp; + + @SuppressWarnings("deprecation") + private OffsetAndMetadata(OffsetMetadata offsetMetadata) { + this( + offsetMetadata, + DEFAULT_TIMESTAMP, + DEFAULT_TIMESTAMP); + } + + public long offset() { + return offsetMetadata.offset(); + } + + public String metadata() { + return offsetMetadata.metadata(); + } + + @Override + public String toString() { + return String.format( + "[%s,CommitTime %d,ExpirationTime %d]", + offsetMetadata, + commitTimestamp, + expireTimestamp + ); + } + +} diff --git a/src/main/java/io/streamnative/kop/offset/OffsetMetadata.java b/src/main/java/io/streamnative/kop/offset/OffsetMetadata.java new file mode 100644 index 0000000000..d00099ed72 --- /dev/null +++ b/src/main/java/io/streamnative/kop/offset/OffsetMetadata.java @@ -0,0 +1,52 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.offset; + +import lombok.Data; +import lombok.experimental.Accessors; + +/** + * Offset Metadata. + */ +@Data +@Accessors(fluent = true) +public class OffsetMetadata { + + public static final long INVALID_OFFSET = -1L; + public static final String NO_METADATA = ""; + public static final OffsetMetadata INVALID_OFFSET_METADATA = + new OffsetMetadata(INVALID_OFFSET, NO_METADATA); + + private final long offset; + private final String metadata; + + public OffsetMetadata(long offset) { + this(offset, NO_METADATA); + } + + public OffsetMetadata(long offset, String metadata) { + this.offset = offset; + this.metadata = metadata; + } + + @Override + public String toString() { + return String.format( + "OffsetMetadata[%d,%s]", + offset, + metadata != null && metadata.length() > 0 ? metadata : "NO_METADATA" + ); + } + +} diff --git a/src/main/java/io/streamnative/kop/offset/OffsetMetadataAndError.java b/src/main/java/io/streamnative/kop/offset/OffsetMetadataAndError.java new file mode 100644 index 0000000000..44dcdec5dd --- /dev/null +++ b/src/main/java/io/streamnative/kop/offset/OffsetMetadataAndError.java @@ -0,0 +1,89 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +package io.streamnative.kop.offset; + +import lombok.AccessLevel; +import lombok.Data; +import lombok.RequiredArgsConstructor; +import org.apache.kafka.common.protocol.Errors; + +/** + * Offset metadata and errors. + */ +@Data +@RequiredArgsConstructor(access = AccessLevel.PRIVATE) +public class OffsetMetadataAndError { + + public static final OffsetMetadataAndError NO_OFFSET = + new OffsetMetadataAndError(OffsetMetadata.INVALID_OFFSET_METADATA, Errors.NONE); + public static final OffsetMetadataAndError GROUP_LOADING = + new OffsetMetadataAndError(OffsetMetadata.INVALID_OFFSET_METADATA, Errors.COORDINATOR_LOAD_IN_PROGRESS); + public static final OffsetMetadataAndError UNKNOWN_MEMBER = + new OffsetMetadataAndError(OffsetMetadata.INVALID_OFFSET_METADATA, Errors.UNKNOWN_MEMBER_ID); + public static final OffsetMetadataAndError NOT_COORDINATOR_FOR_GROUP = + new OffsetMetadataAndError(OffsetMetadata.INVALID_OFFSET_METADATA, Errors.NOT_COORDINATOR); + public static final OffsetMetadataAndError GROUP_COORDINATOR_NOT_AVAILABLE = + new OffsetMetadataAndError(OffsetMetadata.INVALID_OFFSET_METADATA, Errors.COORDINATOR_NOT_AVAILABLE); + public static final OffsetMetadataAndError UNKNOWN_TOPIC_OR_PARTITION = + new OffsetMetadataAndError(OffsetMetadata.INVALID_OFFSET_METADATA, Errors.UNKNOWN_TOPIC_OR_PARTITION); + public static final OffsetMetadataAndError ILLEGAL_GROUP_GENERATION_ID = + new OffsetMetadataAndError(OffsetMetadata.INVALID_OFFSET_METADATA, Errors.ILLEGAL_GENERATION); + + public static OffsetMetadataAndError apply(long offset) { + return new OffsetMetadataAndError( + new OffsetMetadata(offset, OffsetMetadata.NO_METADATA), + Errors.NONE + ); + } + + public static OffsetMetadataAndError apply(Errors errors) { + return new OffsetMetadataAndError( + OffsetMetadata.INVALID_OFFSET_METADATA, + errors + ); + } + + public static OffsetMetadataAndError apply(long offset, + String metadata, + Errors errors) { + return new OffsetMetadataAndError( + new OffsetMetadata(offset, metadata), + errors + ); + } + + private final OffsetMetadata offsetMetadata; + private final Errors error; + + private OffsetMetadataAndError(OffsetMetadata offsetMetadata) { + this(offsetMetadata, Errors.NONE); + } + + public long offset() { + return offsetMetadata.offset(); + } + + public String metadata() { + return offsetMetadata.metadata(); + } + + @Override + public String toString() { + return String.format( + "[%s, Error=%s]", + offsetMetadata, + error + ); + } +} diff --git a/src/main/java/io/streamnative/kop/offset/package-info.java b/src/main/java/io/streamnative/kop/offset/package-info.java new file mode 100644 index 0000000000..b091d4390f --- /dev/null +++ b/src/main/java/io/streamnative/kop/offset/package-info.java @@ -0,0 +1,17 @@ +/** + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +/** + * Kafka Offset related classes. + */ +package io.streamnative.kop.offset; \ No newline at end of file diff --git a/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataManagerTest.java b/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataManagerTest.java index 3b2e4b3911..4880ed272b 100644 --- a/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataManagerTest.java +++ b/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataManagerTest.java @@ -64,6 +64,7 @@ public class GroupMetadataManagerTest extends MockKafkaServiceBaseTest { GroupMetadataManager groupMetadataManager = null; Producer producer = null; Reader consumer = null; + OffsetConfig offsetConfig = OffsetConfig.builder().build(); @Before @Override @@ -83,6 +84,7 @@ public void setup() throws Exception { time = new MockTime(); groupMetadataManager = new GroupMetadataManager( 1, + offsetConfig, producer, consumer, time @@ -154,6 +156,7 @@ void runGroupMetadataManagerTester(final String topicName, .create(); groupMetadataManager = new GroupMetadataManager( 1, + offsetConfig, producer, reader, time diff --git a/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataTest.java b/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataTest.java index adb727f3df..145e5de8f8 100644 --- a/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataTest.java +++ b/src/test/java/io/streamnative/kop/coordinator/group/GroupMetadataTest.java @@ -25,11 +25,16 @@ import static org.junit.Assert.fail; import com.google.common.collect.Sets; +import io.streamnative.kop.coordinator.group.GroupMetadata.CommitRecordMetadataAndOffset; +import io.streamnative.kop.offset.OffsetAndMetadata; +import java.util.HashMap; import java.util.LinkedHashMap; import java.util.Map; +import java.util.Optional; import java.util.Set; import java.util.concurrent.CompletableFuture; import lombok.val; +import org.apache.kafka.common.TopicPartition; import org.junit.Before; import org.junit.Test; @@ -377,6 +382,216 @@ public void testInitNextGenerationEmptyGroup() { assertNull(group.protocolOrNull()); } + @Test + public void testOffsetCommit() { + TopicPartition partition = new TopicPartition("foo", 0); + OffsetAndMetadata offset = OffsetAndMetadata.apply(37); + long commitRecordOffset = 3; + + Map offsets = new HashMap<>(); + offsets.put(partition, offset); + group.prepareOffsetCommit(offsets); + assertTrue(group.hasOffsets()); + assertEquals(Optional.empty(), group.offset(partition)); + + group.onOffsetCommitAppend( + partition, + new CommitRecordMetadataAndOffset(Optional.of(commitRecordOffset), offset)); + assertTrue(group.hasOffsets()); + assertEquals(Optional.of(offset), group.offset(partition)); + } + + @Test + public void testOffsetCommitFailure() { + TopicPartition partition = new TopicPartition("foo", 0); + OffsetAndMetadata offset = OffsetAndMetadata.apply(37); + + Map offsets = new HashMap<>(); + offsets.put(partition, offset); + group.prepareOffsetCommit(offsets); + assertTrue(group.hasOffsets()); + assertEquals(Optional.empty(), group.offset(partition)); + + group.failPendingOffsetWrite(partition, offset); + assertFalse(group.hasOffsets()); + assertEquals(Optional.empty(), group.offset(partition)); + } + + @Test + public void testOffsetCommitFailureWithAnotherPending() { + TopicPartition partition = new TopicPartition("foo", 0); + OffsetAndMetadata firstOffset = OffsetAndMetadata.apply(37); + OffsetAndMetadata secondOffset = OffsetAndMetadata.apply(57); + + Map offsets = new HashMap<>(); + offsets.put(partition, firstOffset); + group.prepareOffsetCommit(offsets); + assertTrue(group.hasOffsets()); + assertEquals(Optional.empty(), group.offset(partition)); + + offsets = new HashMap<>(); + offsets.put(partition, secondOffset); + group.prepareOffsetCommit(offsets); + assertTrue(group.hasOffsets()); + + group.failPendingOffsetWrite(partition, firstOffset); + assertTrue(group.hasOffsets()); + assertEquals(Optional.empty(), group.offset(partition)); + + group.onOffsetCommitAppend(partition, new CommitRecordMetadataAndOffset(Optional.of(3L), secondOffset)); + assertTrue(group.hasOffsets()); + assertEquals(Optional.of(secondOffset), group.offset(partition)); + } + + @Test + public void testOffsetCommitWithAnotherPending() { + TopicPartition partition = new TopicPartition("foo", 0); + OffsetAndMetadata firstOffset = OffsetAndMetadata.apply(37); + OffsetAndMetadata secondOffset = OffsetAndMetadata.apply(57); + + Map offsets = new HashMap<>(); + offsets.put(partition, firstOffset); + group.prepareOffsetCommit(offsets); + assertTrue(group.hasOffsets()); + assertEquals(Optional.empty(), group.offset(partition)); + + offsets = new HashMap<>(); + offsets.put(partition, secondOffset); + group.prepareOffsetCommit(offsets); + assertTrue(group.hasOffsets()); + + group.onOffsetCommitAppend(partition, new CommitRecordMetadataAndOffset(Optional.of(4L), firstOffset)); + assertTrue(group.hasOffsets()); + assertEquals(Optional.of(firstOffset), group.offset(partition)); + + group.onOffsetCommitAppend(partition, new CommitRecordMetadataAndOffset(Optional.of(5L), secondOffset)); + assertTrue(group.hasOffsets()); + assertEquals(Optional.of(secondOffset), group.offset(partition)); + } + + @Test + public void testConsumerBeatsTransactionalOffsetCommit() { + TopicPartition partition = new TopicPartition("foo", 0); + long producerId = 13232L; + OffsetAndMetadata txnOffsetCommit = OffsetAndMetadata.apply(37); + OffsetAndMetadata consumerOffsetCommit = OffsetAndMetadata.apply(57); + + Map offsets = new HashMap<>(); + offsets.put(partition, txnOffsetCommit); + group.prepareTxnOffsetCommit(producerId, offsets); + assertTrue(group.hasOffsets()); + assertEquals(Optional.empty(), group.offset(partition)); + + offsets = new HashMap<>(); + offsets.put(partition, consumerOffsetCommit); + group.prepareOffsetCommit(offsets); + assertTrue(group.hasOffsets()); + + group.onTxnOffsetCommitAppend(producerId, partition, + new CommitRecordMetadataAndOffset(Optional.of(3L), txnOffsetCommit)); + group.onOffsetCommitAppend(partition, + new CommitRecordMetadataAndOffset(Optional.of(4L), consumerOffsetCommit)); + assertTrue(group.hasOffsets()); + assertEquals(Optional.of(consumerOffsetCommit), group.offset(partition)); + + group.completePendingTxnOffsetCommit(producerId, true); + assertTrue(group.hasOffsets()); + // This is the crucial assertion which validates that we materialize offsets in offset order, + // not transactional order. + assertEquals(Optional.of(consumerOffsetCommit), group.offset(partition)); + } + + @Test + public void testTransactionBeatsConsumerOffsetCommit() { + TopicPartition partition = new TopicPartition("foo", 0); + long producerId = 13232L; + OffsetAndMetadata txnOffsetCommit = OffsetAndMetadata.apply(37); + OffsetAndMetadata consumerOffsetCommit = OffsetAndMetadata.apply(57); + + Map offsets = new HashMap<>(); + offsets.put(partition, txnOffsetCommit); + group.prepareTxnOffsetCommit(producerId, offsets); + assertTrue(group.hasOffsets()); + assertEquals(Optional.empty(), group.offset(partition)); + + offsets = new HashMap<>(); + offsets.put(partition, consumerOffsetCommit); + group.prepareOffsetCommit(offsets); + assertTrue(group.hasOffsets()); + + group.onOffsetCommitAppend( + partition, new CommitRecordMetadataAndOffset(Optional.of(3L), consumerOffsetCommit)); + group.onTxnOffsetCommitAppend(producerId, partition, + new CommitRecordMetadataAndOffset(Optional.of(4L), txnOffsetCommit)); + assertTrue(group.hasOffsets()); + // The transactional offset commit hasn't been committed yet, so we should materialize + // the consumer offset commit. + assertEquals(Optional.of(consumerOffsetCommit), group.offset(partition)); + + group.completePendingTxnOffsetCommit(producerId, true); + assertTrue(group.hasOffsets()); + // The transactional offset commit has been materialized and the transactional commit record + // is later in the log, so it should be materialized. + assertEquals(Optional.of(txnOffsetCommit), group.offset(partition)); + } + + @Test + public void testTransactionalCommitIsAbortedAndConsumerCommitWins() { + TopicPartition partition = new TopicPartition("foo", 0); + long producerId = 13232L; + OffsetAndMetadata txnOffsetCommit = OffsetAndMetadata.apply(37); + OffsetAndMetadata consumerOffsetCommit = OffsetAndMetadata.apply(57); + + Map offsets = new HashMap<>(); + offsets.put(partition, txnOffsetCommit); + group.prepareTxnOffsetCommit(producerId, offsets); + assertTrue(group.hasOffsets()); + assertEquals(Optional.empty(), group.offset(partition)); + + offsets = new HashMap<>(); + offsets.put(partition, consumerOffsetCommit); + group.prepareOffsetCommit(offsets); + assertTrue(group.hasOffsets()); + + group.onOffsetCommitAppend(partition, + new CommitRecordMetadataAndOffset(Optional.of(3L), consumerOffsetCommit)); + group.onTxnOffsetCommitAppend(producerId, partition, + new CommitRecordMetadataAndOffset(Optional.of(4L), txnOffsetCommit)); + assertTrue(group.hasOffsets()); + // The transactional offset commit hasn't been committed yet, so we should materialize the consumer + // offset commit. + assertEquals(Optional.of(consumerOffsetCommit), group.offset(partition)); + + group.completePendingTxnOffsetCommit(producerId, false); + assertTrue(group.hasOffsets()); + // The transactional offset commit should be discarded and the consumer offset commit should continue to be + // materialized. + assertFalse(group.hasPendingOffsetCommitsFromProducer(producerId)); + assertEquals(Optional.of(consumerOffsetCommit), group.offset(partition)); + } + + @Test + public void testFailedTxnOffsetCommitLeavesNoPendingState() { + TopicPartition partition = new TopicPartition("foo", 0); + long producerId = 13232L; + OffsetAndMetadata txnOffsetCommit = OffsetAndMetadata.apply(37); + + Map offsets = new HashMap<>(); + offsets.put(partition, txnOffsetCommit); + group.prepareTxnOffsetCommit(producerId, offsets); + assertTrue(group.hasPendingOffsetCommitsFromProducer(producerId)); + assertTrue(group.hasOffsets()); + assertEquals(Optional.empty(), group.offset(partition)); + group.failPendingTxnOffsetCommit(producerId, partition); + assertFalse(group.hasOffsets()); + assertFalse(group.hasPendingOffsetCommitsFromProducer(producerId)); + + // The commit marker should now have no effect. + group.completePendingTxnOffsetCommit(producerId, true); + assertFalse(group.hasOffsets()); + assertFalse(group.hasPendingOffsetCommitsFromProducer(producerId)); + } + private void assertState(GroupMetadata group, GroupState targetState) { Set states = Sets.newHashSet( Stable, PreparingRebalance, CompletingRebalance, Dead @@ -387,5 +602,4 @@ private void assertState(GroupMetadata group, GroupState targetState) { assertTrue(group.is(targetState)); } - }