Viacheslav Klimov
2 years 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