Viacheslav Klimov
1 year ago
committed by
GitHub
GitHub
29 changed files with 483 additions and 323 deletions
@ -1,98 +0,0 @@ |
|||
/** |
|||
* Copyright © 2016-2025 The Thingsboard Authors |
|||
* |
|||
* 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 org.thingsboard.server.queue.common.consumer; |
|||
|
|||
import lombok.Getter; |
|||
import lombok.extern.slf4j.Slf4j; |
|||
import org.thingsboard.server.common.msg.queue.TopicPartitionInfo; |
|||
import org.thingsboard.server.queue.TbQueueMsg; |
|||
|
|||
import java.util.Collections; |
|||
import java.util.HashSet; |
|||
import java.util.Set; |
|||
import java.util.concurrent.ConcurrentHashMap; |
|||
import java.util.concurrent.locks.ReadWriteLock; |
|||
import java.util.concurrent.locks.ReentrantReadWriteLock; |
|||
|
|||
import static org.thingsboard.server.common.msg.queue.TopicPartitionInfo.withTopic; |
|||
|
|||
@Slf4j |
|||
public class QueueStateService<E extends TbQueueMsg, S extends TbQueueMsg> { |
|||
|
|||
private PartitionedQueueConsumerManager<S> stateConsumer; |
|||
private PartitionedQueueConsumerManager<E> eventConsumer; |
|||
|
|||
@Getter |
|||
private Set<TopicPartitionInfo> partitions; |
|||
private final Set<TopicPartitionInfo> partitionsInProgress = ConcurrentHashMap.newKeySet(); |
|||
private boolean initialized; |
|||
|
|||
private final ReadWriteLock partitionsLock = new ReentrantReadWriteLock(); |
|||
|
|||
public void init(PartitionedQueueConsumerManager<S> stateConsumer, PartitionedQueueConsumerManager<E> eventConsumer) { |
|||
this.stateConsumer = stateConsumer; |
|||
this.eventConsumer = eventConsumer; |
|||
} |
|||
|
|||
public void update(Set<TopicPartitionInfo> newPartitions) { |
|||
newPartitions = withTopic(newPartitions, stateConsumer.getTopic()); |
|||
var writeLock = partitionsLock.writeLock(); |
|||
writeLock.lock(); |
|||
Set<TopicPartitionInfo> oldPartitions = this.partitions != null ? this.partitions : Collections.emptySet(); |
|||
Set<TopicPartitionInfo> addedPartitions; |
|||
Set<TopicPartitionInfo> removedPartitions; |
|||
try { |
|||
addedPartitions = new HashSet<>(newPartitions); |
|||
addedPartitions.removeAll(oldPartitions); |
|||
removedPartitions = new HashSet<>(oldPartitions); |
|||
removedPartitions.removeAll(newPartitions); |
|||
this.partitions = newPartitions; |
|||
} finally { |
|||
writeLock.unlock(); |
|||
} |
|||
|
|||
if (!removedPartitions.isEmpty()) { |
|||
stateConsumer.removePartitions(removedPartitions); |
|||
eventConsumer.removePartitions(withTopic(removedPartitions, eventConsumer.getTopic())); |
|||
} |
|||
|
|||
if (!addedPartitions.isEmpty()) { |
|||
partitionsInProgress.addAll(addedPartitions); |
|||
stateConsumer.addPartitions(addedPartitions, partition -> { |
|||
var readLock = partitionsLock.readLock(); |
|||
readLock.lock(); |
|||
try { |
|||
partitionsInProgress.remove(partition); |
|||
log.info("Finished partition {} (still in progress: {})", partition, partitionsInProgress); |
|||
if (partitionsInProgress.isEmpty()) { |
|||
log.info("All partitions processed"); |
|||
} |
|||
if (this.partitions.contains(partition)) { |
|||
eventConsumer.addPartitions(Set.of(partition.withTopic(eventConsumer.getTopic()))); |
|||
} |
|||
} finally { |
|||
readLock.unlock(); |
|||
} |
|||
}); |
|||
} |
|||
initialized = true; |
|||
} |
|||
|
|||
public Set<TopicPartitionInfo> getPartitionsInProgress() { |
|||
return initialized ? partitionsInProgress : null; |
|||
} |
|||
|
|||
} |
|||
@ -0,0 +1,27 @@ |
|||
/** |
|||
* Copyright © 2016-2025 The Thingsboard Authors |
|||
* |
|||
* 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 org.thingsboard.server.queue.common.state; |
|||
|
|||
import org.thingsboard.server.queue.TbQueueMsg; |
|||
import org.thingsboard.server.queue.common.consumer.PartitionedQueueConsumerManager; |
|||
|
|||
public class DefaultQueueStateService<E extends TbQueueMsg, S extends TbQueueMsg> extends QueueStateService<E, S> { |
|||
|
|||
public DefaultQueueStateService(PartitionedQueueConsumerManager<E> eventConsumer) { |
|||
super(eventConsumer); |
|||
} |
|||
|
|||
} |
|||
@ -0,0 +1,81 @@ |
|||
/** |
|||
* Copyright © 2016-2025 The Thingsboard Authors |
|||
* |
|||
* 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 org.thingsboard.server.queue.common.state; |
|||
|
|||
import lombok.extern.slf4j.Slf4j; |
|||
import org.thingsboard.server.common.msg.queue.TopicPartitionInfo; |
|||
import org.thingsboard.server.queue.TbQueueMsg; |
|||
import org.thingsboard.server.queue.common.consumer.PartitionedQueueConsumerManager; |
|||
import org.thingsboard.server.queue.discovery.QueueKey; |
|||
|
|||
import java.util.Set; |
|||
|
|||
import static org.thingsboard.server.common.msg.queue.TopicPartitionInfo.withTopic; |
|||
|
|||
@Slf4j |
|||
public class KafkaQueueStateService<E extends TbQueueMsg, S extends TbQueueMsg> extends QueueStateService<E, S> { |
|||
|
|||
private final PartitionedQueueConsumerManager<S> stateConsumer; |
|||
|
|||
public KafkaQueueStateService(PartitionedQueueConsumerManager<E> eventConsumer, PartitionedQueueConsumerManager<S> stateConsumer) { |
|||
super(eventConsumer); |
|||
this.stateConsumer = stateConsumer; |
|||
} |
|||
|
|||
@Override |
|||
protected void addPartitions(QueueKey queueKey, Set<TopicPartitionInfo> partitions) { |
|||
Set<TopicPartitionInfo> statePartitions = withTopic(partitions, stateConsumer.getTopic()); |
|||
partitionsInProgress.addAll(statePartitions); |
|||
stateConsumer.addPartitions(statePartitions, statePartition -> { |
|||
var readLock = partitionsLock.readLock(); |
|||
readLock.lock(); |
|||
try { |
|||
partitionsInProgress.remove(statePartition); |
|||
log.info("Finished partition {} (still in progress: {})", statePartition, partitionsInProgress); |
|||
if (partitionsInProgress.isEmpty()) { |
|||
log.info("All partitions processed"); |
|||
} |
|||
|
|||
TopicPartitionInfo eventPartition = statePartition.withTopic(eventConsumer.getTopic()); |
|||
if (this.partitions.get(queueKey).contains(eventPartition)) { |
|||
eventConsumer.addPartitions(Set.of(eventPartition)); |
|||
} |
|||
} finally { |
|||
readLock.unlock(); |
|||
} |
|||
}); |
|||
} |
|||
|
|||
@Override |
|||
protected void removePartitions(QueueKey queueKey, Set<TopicPartitionInfo> partitions) { |
|||
super.removePartitions(queueKey, partitions); |
|||
stateConsumer.removePartitions(withTopic(partitions, stateConsumer.getTopic())); |
|||
} |
|||
|
|||
@Override |
|||
protected void deletePartitions(Set<TopicPartitionInfo> partitions) { |
|||
super.deletePartitions(partitions); |
|||
stateConsumer.delete(withTopic(partitions, stateConsumer.getTopic())); |
|||
} |
|||
|
|||
@Override |
|||
public void stop() { |
|||
super.stop(); |
|||
stateConsumer.stop(); |
|||
stateConsumer.awaitStop(); |
|||
} |
|||
|
|||
} |
|||
@ -0,0 +1,114 @@ |
|||
/** |
|||
* Copyright © 2016-2025 The Thingsboard Authors |
|||
* |
|||
* 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 org.thingsboard.server.queue.common.state; |
|||
|
|||
import lombok.Getter; |
|||
import lombok.extern.slf4j.Slf4j; |
|||
import org.thingsboard.server.common.msg.queue.TopicPartitionInfo; |
|||
import org.thingsboard.server.queue.TbQueueMsg; |
|||
import org.thingsboard.server.queue.common.consumer.PartitionedQueueConsumerManager; |
|||
import org.thingsboard.server.queue.discovery.QueueKey; |
|||
|
|||
import java.util.Collections; |
|||
import java.util.HashMap; |
|||
import java.util.HashSet; |
|||
import java.util.Map; |
|||
import java.util.Set; |
|||
import java.util.concurrent.ConcurrentHashMap; |
|||
import java.util.concurrent.locks.ReadWriteLock; |
|||
import java.util.concurrent.locks.ReentrantReadWriteLock; |
|||
|
|||
import static org.thingsboard.server.common.msg.queue.TopicPartitionInfo.withTopic; |
|||
|
|||
@Slf4j |
|||
public abstract class QueueStateService<E extends TbQueueMsg, S extends TbQueueMsg> { |
|||
|
|||
protected final PartitionedQueueConsumerManager<E> eventConsumer; |
|||
|
|||
@Getter |
|||
protected final Map<QueueKey, Set<TopicPartitionInfo>> partitions = new HashMap<>(); |
|||
protected final Set<TopicPartitionInfo> partitionsInProgress = ConcurrentHashMap.newKeySet(); |
|||
protected boolean initialized; |
|||
|
|||
protected final ReadWriteLock partitionsLock = new ReentrantReadWriteLock(); |
|||
|
|||
protected QueueStateService(PartitionedQueueConsumerManager<E> eventConsumer) { |
|||
this.eventConsumer = eventConsumer; |
|||
} |
|||
|
|||
public void update(QueueKey queueKey, Set<TopicPartitionInfo> newPartitions) { |
|||
newPartitions = withTopic(newPartitions, eventConsumer.getTopic()); |
|||
var writeLock = partitionsLock.writeLock(); |
|||
writeLock.lock(); |
|||
Set<TopicPartitionInfo> oldPartitions = this.partitions.getOrDefault(queueKey, Collections.emptySet()); |
|||
Set<TopicPartitionInfo> addedPartitions; |
|||
Set<TopicPartitionInfo> removedPartitions; |
|||
try { |
|||
addedPartitions = new HashSet<>(newPartitions); |
|||
addedPartitions.removeAll(oldPartitions); |
|||
removedPartitions = new HashSet<>(oldPartitions); |
|||
removedPartitions.removeAll(newPartitions); |
|||
this.partitions.put(queueKey, newPartitions); |
|||
} finally { |
|||
writeLock.unlock(); |
|||
} |
|||
|
|||
if (!removedPartitions.isEmpty()) { |
|||
removePartitions(queueKey, removedPartitions); |
|||
} |
|||
|
|||
if (!addedPartitions.isEmpty()) { |
|||
addPartitions(queueKey, addedPartitions); |
|||
} |
|||
initialized = true; |
|||
} |
|||
|
|||
protected void addPartitions(QueueKey queueKey, Set<TopicPartitionInfo> partitions) { |
|||
eventConsumer.addPartitions(partitions); |
|||
} |
|||
|
|||
protected void removePartitions(QueueKey queueKey, Set<TopicPartitionInfo> partitions) { |
|||
eventConsumer.removePartitions(partitions); |
|||
} |
|||
|
|||
public void delete(Set<TopicPartitionInfo> partitions) { |
|||
if (partitions.isEmpty()) { |
|||
return; |
|||
} |
|||
var writeLock = partitionsLock.writeLock(); |
|||
writeLock.lock(); |
|||
try { |
|||
this.partitions.values().forEach(tpis -> tpis.removeAll(partitions)); |
|||
} finally { |
|||
writeLock.unlock(); |
|||
} |
|||
deletePartitions(partitions); |
|||
} |
|||
|
|||
protected void deletePartitions(Set<TopicPartitionInfo> partitions) { |
|||
eventConsumer.delete(withTopic(partitions, eventConsumer.getTopic())); |
|||
} |
|||
|
|||
public Set<TopicPartitionInfo> getPartitionsInProgress() { |
|||
return initialized ? partitionsInProgress : null; |
|||
} |
|||
|
|||
public void stop() { |
|||
eventConsumer.stop(); |
|||
eventConsumer.awaitStop(); |
|||
} |
|||
|
|||
} |
|||
Loading…
Reference in new issue