@ -13,7 +13,7 @@
* See the License for the specific language governing permissions and
* limitations under the License .
* /
package org.thingsboard.server.service. queue.consumer ;
package org.thingsboard.server.queue.common .consumer ;
import lombok.Builder ;
import lombok.Getter ;
@ -23,10 +23,9 @@ import org.thingsboard.server.common.data.queue.QueueConfig;
import org.thingsboard.server.common.msg.queue.TopicPartitionInfo ;
import org.thingsboard.server.queue.TbQueueConsumer ;
import org.thingsboard.server.queue.TbQueueMsg ;
import org.thingsboard.server.queue.common.consumer.TbQueueConsumerManagerTask.UpdateConfigTask ;
import org.thingsboard.server.queue.common.consumer.TbQueueConsumerManagerTask.UpdatePartitionsTask ;
import org.thingsboard.server.queue.discovery.QueueKey ;
import org.thingsboard.server.service.queue.ruleengine.QueueEvent ;
import org.thingsboard.server.service.queue.ruleengine.TbQueueConsumerManagerTask ;
import org.thingsboard.server.service.queue.ruleengine.TbQueueConsumerTask ;
import java.util.Collection ;
import java.util.Collections ;
@ -34,6 +33,7 @@ import java.util.HashMap;
import java.util.HashSet ;
import java.util.List ;
import java.util.Map ;
import java.util.Optional ;
import java.util.Set ;
import java.util.concurrent.ConcurrentLinkedQueue ;
import java.util.concurrent.ExecutorService ;
@ -42,6 +42,7 @@ import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.TimeUnit ;
import java.util.concurrent.locks.ReentrantLock ;
import java.util.function.BiFunction ;
import java.util.function.Consumer ;
import java.util.stream.Collectors ;
@Slf4j
@ -52,8 +53,11 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
protected C config ;
protected final MsgPackProcessor < M , C > msgPackProcessor ;
protected final BiFunction < C , Integer , TbQueueConsumer < M > > consumerCreator ;
@Getter
protected final ExecutorService consumerExecutor ;
@Getter
protected final ScheduledExecutorService scheduler ;
@Getter
protected final ExecutorService taskExecutor ;
private final java . util . Queue < TbQueueConsumerManagerTask > tasks = new ConcurrentLinkedQueue < > ( ) ;
@ -85,20 +89,24 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
public void init ( C config ) {
this . config = config ;
this . consumerWrapper = createConsumerWrapper ( config ) ;
log . debug ( "[{}] Initialized consumer for queue: {}" , queueKey , config ) ;
}
protected ConsumerWrapper < M > createConsumerWrapper ( C config ) {
if ( config . isConsumerPerPartition ( ) ) {
this . consumerWrapper = new ConsumerPerPartitionWrapper ( ) ;
retu rn new ConsumerPerPartitionWrapper ( ) ;
} else {
this . consumerWrapper = new SingleConsumerWrapper ( ) ;
retu rn new SingleConsumerWrapper ( ) ;
}
log . debug ( "[{}] Initialized consumer for queue: {}" , queueKey , config ) ;
}
public void update ( C config ) {
addTask ( TbQueueConsumerManagerTask . configUpdate ( config ) ) ;
addTask ( new UpdateConfigTask ( config ) ) ;
}
public void update ( Set < TopicPartitionInfo > partitions ) {
addTask ( TbQueueConsumerManagerTask . partitionChange ( partitions ) ) ;
addTask ( new UpdatePartitionsTask ( partitions ) ) ;
}
protected void addTask ( TbQueueConsumerManagerTask todo ) {
@ -123,10 +131,10 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
}
log . trace ( "[{}] Processing task: {}" , queueKey , task ) ;
if ( task . getEvent ( ) = = QueueEvent . PARTITION_CHANGE ) {
newPartitions = task . getP artitions( ) ;
} else if ( task . getEvent ( ) = = QueueEvent . CONFIG_UPDATE ) {
newConfig = ( C ) task . getC onfig( ) ;
if ( task instanceof UpdatePartitionsTask updatePartitionsTask ) {
newPartitions = upda tePartitionsT ask. p artitions( ) ;
} else if ( task instanceof UpdateConfigTask updateConfigTask ) {
newConfig = ( C ) upda teConfigT ask. c onfig( ) ;
} else {
processTask ( task ) ;
}
@ -182,7 +190,7 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
}
}
public void doUpdate ( Set < TopicPartitionInfo > partitions ) {
private void doUpdate ( Set < TopicPartitionInfo > partitions ) {
this . partitions = partitions ;
consumerWrapper . updatePartitions ( partitions ) ;
}
@ -197,6 +205,15 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
log . error ( "Failure in consumer loop" , e ) ;
}
log . info ( "[{}] Consumer stopped" , consumerTask . getKey ( ) ) ;
try {
Runnable callback = consumerTask . getCallback ( ) ;
if ( callback ! = null ) {
callback . run ( ) ;
}
} catch ( Throwable t ) {
log . error ( "Failed to execute finish callback" , t ) ;
}
} ) ;
consumerTask . setTask ( consumerLoop ) ;
}
@ -241,13 +258,13 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
awaitStop ( 30 ) ;
}
public void awaitStop ( long timeoutSec ) {
private void awaitStop ( int timeoutSec ) {
log . debug ( "[{}] Waiting for consumers to stop" , queueKey ) ;
consumerWrapper . getConsumers ( ) . forEach ( consumerTask - > consumerTask . awaitCompletion ( timeoutSec ) ) ;
log . debug ( "[{}] Unsubscribed and stopped consumers" , queueKey ) ;
}
private static String partitionsToString ( Collection < TopicPartitionInfo > partitions ) {
static String partitionsToString ( Collection < TopicPartitionInfo > partitions ) {
return partitions . stream ( ) . map ( TopicPartitionInfo : : getFullTopicName ) . collect ( Collectors . joining ( ", " , "[" , "]" ) ) ;
}
@ -273,15 +290,24 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
Set < TopicPartitionInfo > removedPartitions = new HashSet < > ( consumers . keySet ( ) ) ;
removedPartitions . removeAll ( partitions ) ;
log . info ( "[{}] Added partitions: {}, removed partitions: {}" , queueKey , partitionsToString ( addedPartitions ) , partitionsToString ( removedPartitions ) ) ;
removePartitions ( removedPartitions ) ;
addPartitions ( addedPartitions , null ) ;
}
removedPartitions . forEach ( ( tpi ) - > consumers . get ( tpi ) . initiateStop ( ) ) ;
removedPartitions . forEach ( ( tpi ) - > consumers . remove ( tpi ) . awaitCompletion ( ) ) ;
protected void removePartitions ( Set < TopicPartitionInfo > removedPartitions ) {
removedPartitions . forEach ( ( tpi ) - > Optional . ofNullable ( consumers . get ( tpi ) ) . ifPresent ( TbQueueConsumerTask : : initiateStop ) ) ;
removedPartitions . forEach ( ( tpi ) - > Optional . ofNullable ( consumers . remove ( tpi ) ) . ifPresent ( TbQueueConsumerTask : : awaitCompletion ) ) ;
}
addedPartitions . forEach ( ( tpi ) - > {
protected void addPartitions ( Set < TopicPartitionInfo > partitions , Consumer < TopicPartitionInfo > onStop ) {
partitions . forEach ( tpi - > {
Integer partitionId = tpi . getPartition ( ) . orElse ( - 1 ) ;
String key = queueKey + "-" + partitionId ;
TbQueueConsumerTask < M > consumer = new TbQueueConsumerTask < > ( key , ( ) - > consumerCreator . apply ( config , partitionId ) ) ;
Runnable callback = onStop ! = null ? ( ) - > onStop . accept ( tpi ) : null ;
TbQueueConsumerTask < M > consumer = new TbQueueConsumerTask < > ( key , ( ) - > consumerCreator . apply ( config , partitionId ) , callback ) ;
consumers . put ( tpi , consumer ) ;
consumer . subscribe ( Set . of ( tpi ) ) ;
launchConsumer ( consumer ) ;
@ -310,7 +336,7 @@ public class MainQueueConsumerManager<M extends TbQueueMsg, C extends QueueConfi
}
if ( consumer = = null ) {
consumer = new TbQueueConsumerTask < > ( queueKey , ( ) - > consumerCreator . apply ( config , null ) ) ; // no partitionId passed
consumer = new TbQueueConsumerTask < > ( queueKey , ( ) - > consumerCreator . apply ( config , null ) , null ) ; // no partitionId passed
}
consumer . subscribe ( partitions ) ;
if ( ! consumer . isRunning ( ) ) {
Andrii Shvaika
2 years ago