@ -18,6 +18,7 @@ package org.thingsboard.server.queue.discovery;
import com.datastax.oss.driver.api.core.uuid.Uuids ;
import lombok.SneakyThrows ;
import lombok.extern.slf4j.Slf4j ;
import org.apache.commons.collections4.ListUtils ;
import org.junit.Assert ;
import org.junit.Before ;
import org.junit.Test ;
@ -25,24 +26,35 @@ import org.junit.runner.RunWith;
import org.mockito.junit.MockitoJUnitRunner ;
import org.springframework.context.ApplicationEventPublisher ;
import org.springframework.test.util.ReflectionTestUtils ;
import org.thingsboard.server.common.data.DataConstants ;
import org.thingsboard.server.common.data.id.DeviceId ;
import org.thingsboard.server.common.data.id.QueueId ;
import org.thingsboard.server.common.data.id.TenantId ;
import org.thingsboard.server.common.data.id.TenantProfileId ;
import org.thingsboard.server.common.data.id.UUIDBased ;
import org.thingsboard.server.common.data.queue.Queue ;
import org.thingsboard.server.common.msg.queue.ServiceType ;
import org.thingsboard.server.common.msg.queue.TopicPartitionInfo ;
import org.thingsboard.server.gen.transport.TransportProtos ;
import org.thingsboard.server.gen.transport.TransportProtos.ServiceInfo ;
import java.text.SimpleDateFormat ;
import java.util.ArrayList ;
import java.util.Collection ;
import java.util.Collections ;
import java.util.Comparator ;
import java.util.HashMap ;
import java.util.List ;
import java.util.Map ;
import java.util.Random ;
import java.util.UUID ;
import java.util.concurrent.TimeUnit ;
import java.util.stream.Collectors ;
import java.util.stream.Stream ;
import static org.assertj.core.api.Assertions.assertThat ;
import static org.mockito.ArgumentMatchers.eq ;
import static org.mockito.Mockito.mock ;
import static org.mockito.Mockito.when ;
@Slf4j
@RunWith ( MockitoJUnitRunner . class )
@ -57,7 +69,7 @@ public class HashPartitionServiceTest {
private ApplicationEventPublisher applicationEventPublisher ;
private QueueRoutingInfoService queueRoutingInfoService ;
private String hashFunctionName = "sha256 " ;
private String hashFunctionName = "murmur3_128 " ;
@Before
public void setup ( ) throws Exception {
@ -74,15 +86,15 @@ public class HashPartitionServiceTest {
ReflectionTestUtils . setField ( clusterRoutingService , "vcTopic" , "tb.vc" ) ;
ReflectionTestUtils . setField ( clusterRoutingService , "vcPartitions" , 10 ) ;
ReflectionTestUtils . setField ( clusterRoutingService , "hashFunctionName" , hashFunctionName ) ;
TransportProtos . ServiceInfo currentServer = TransportProtos . ServiceInfo . newBuilder ( )
ServiceInfo currentServer = ServiceInfo . newBuilder ( )
. setServiceId ( "tb-core-0" )
. addAllServiceTypes ( Collections . singletonList ( ServiceType . TB_CORE . name ( ) ) )
. build ( ) ;
// when(queueService.resolve(Mockito.any(), Mockito.anyString())).thenAnswer(i -> i.getArguments()[1]);
// when(discoveryService.getServiceInfo()).thenReturn(currentServer);
List < TransportProtos . ServiceInfo > otherServers = new ArrayList < > ( ) ;
List < ServiceInfo > otherServers = new ArrayList < > ( ) ;
for ( int i = 1 ; i < SERVER_COUNT ; i + + ) {
otherServers . add ( TransportProtos . ServiceInfo . newBuilder ( )
otherServers . add ( ServiceInfo . newBuilder ( )
. setServiceId ( "tb-rule-" + i )
. addAllServiceTypes ( Collections . singletonList ( ServiceType . TB_CORE . name ( ) ) )
. build ( ) ) ;
@ -122,10 +134,10 @@ public class HashPartitionServiceTest {
int queueCount = 3 ;
int partitionCount = 3 ;
List < TransportProtos . ServiceInfo > services = new ArrayList < > ( ) ;
List < ServiceInfo > services = new ArrayList < > ( ) ;
for ( int i = 0 ; i < serverCount ; i + + ) {
services . add ( TransportProtos . ServiceInfo . newBuilder ( ) . setServiceId ( "RE-" + i ) . build ( ) ) ;
services . add ( ServiceInfo . newBuilder ( ) . setServiceId ( "RE-" + i ) . build ( ) ) ;
}
long start = System . currentTimeMillis ( ) ;
@ -140,7 +152,7 @@ public class HashPartitionServiceTest {
for ( int queueIndex = 0 ; queueIndex < queueCount ; queueIndex + + ) {
QueueKey queueKey = new QueueKey ( ServiceType . TB_RULE_ENGINE , "queue" + queueIndex , tenantId ) ;
for ( int partition = 0 ; partition < partitionCount ; partition + + ) {
TransportProtos . ServiceInfo serviceInfo = clusterRoutingService . resolveByPartitionIdx ( services , queueKey , partition ) ;
ServiceInfo serviceInfo = clusterRoutingService . resolveByPartitionIdx ( services , queueKey , partition ) ;
String serviceId = serviceInfo . getServiceId ( ) ;
map . put ( serviceId , map . get ( serviceId ) + 1 ) ;
}
@ -163,4 +175,103 @@ public class HashPartitionServiceTest {
Assert . assertTrue ( diffPercent < maxDiffPercent ) ;
}
@Test
public void testPartitionsAssignmentWithDedicatedServers ( ) {
int isolatedProfilesCount = 5 ;
int tenantsCountPerProfile = 100 ;
int dedicatedServerSetsCount = 3 ;
int serversCountPerSet = 3 ;
int profilesPerSet = ( int ) Math . ceil ( ( double ) isolatedProfilesCount / dedicatedServerSetsCount ) ;
List < TenantProfileId > isolatedTenantProfiles = Stream . generate ( ( ) - > new TenantProfileId ( UUID . randomUUID ( ) ) )
. limit ( isolatedProfilesCount ) . collect ( Collectors . toList ( ) ) ;
Map < TenantId , TenantProfileId > tenants = new HashMap < > ( ) ;
for ( TenantProfileId tenantProfileId : isolatedTenantProfiles ) {
for ( int i = 0 ; i < tenantsCountPerProfile ; i + + ) {
tenants . put ( new TenantId ( UUID . randomUUID ( ) ) , tenantProfileId ) ;
}
}
List < Queue > queues = new ArrayList < > ( ) ;
Queue systemQueue = new Queue ( ) ;
systemQueue . setTenantId ( TenantId . SYS_TENANT_ID ) ;
systemQueue . setName ( "Main" ) ;
systemQueue . setTopic ( DataConstants . MAIN_QUEUE_TOPIC ) ;
systemQueue . setPartitions ( 10 ) ;
systemQueue . setId ( new QueueId ( UUID . randomUUID ( ) ) ) ;
queues . add ( systemQueue ) ;
tenants . forEach ( ( tenantId , profileId ) - > {
Queue isolatedQueue = new Queue ( ) ;
isolatedQueue . setTenantId ( tenantId ) ;
isolatedQueue . setName ( "Main" ) ;
isolatedQueue . setTopic ( DataConstants . MAIN_QUEUE_TOPIC ) ;
isolatedQueue . setPartitions ( 2 ) ;
isolatedQueue . setId ( new QueueId ( UUID . randomUUID ( ) ) ) ;
queues . add ( isolatedQueue ) ;
when ( routingInfoService . getRoutingInfo ( eq ( tenantId ) ) ) . thenReturn ( new TenantRoutingInfo ( tenantId , profileId , true ) ) ;
} ) ;
when ( queueRoutingInfoService . getAllQueuesRoutingInfo ( ) ) . thenReturn ( queues . stream ( )
. map ( QueueRoutingInfo : : new ) . collect ( Collectors . toList ( ) ) ) ;
List < ServiceInfo > ruleEngines = new ArrayList < > ( ) ;
Map < TenantProfileId , List < ServiceInfo > > dedicatedServers = new HashMap < > ( ) ;
int serviceId = 0 ;
for ( int i = 0 ; i < serversCountPerSet ; i + + ) {
ServiceInfo commonServer = ServiceInfo . newBuilder ( )
. setServiceId ( "tb-rule-engine-" + serviceId )
. addAllServiceTypes ( List . of ( ServiceType . TB_RULE_ENGINE . name ( ) ) )
. build ( ) ;
ruleEngines . add ( commonServer ) ;
serviceId + + ;
}
for ( int i = 0 ; i < dedicatedServerSetsCount ; i + + ) {
List < TenantProfileId > assignedProfiles = ListUtils . partition ( isolatedTenantProfiles , profilesPerSet ) . get ( i ) ;
for ( int j = 0 ; j < serversCountPerSet ; j + + ) {
ServiceInfo dedicatedServer = ServiceInfo . newBuilder ( )
. setServiceId ( "tb-rule-engine-" + serviceId )
. addAllServiceTypes ( List . of ( ServiceType . TB_RULE_ENGINE . name ( ) ) )
. addAllAssignedTenantProfiles ( assignedProfiles . stream ( ) . map ( UUIDBased : : toString ) . collect ( Collectors . toList ( ) ) )
. build ( ) ;
ruleEngines . add ( dedicatedServer ) ;
serviceId + + ;
for ( TenantProfileId assignedProfileId : assignedProfiles ) {
dedicatedServers . computeIfAbsent ( assignedProfileId , p - > new ArrayList < > ( ) ) . add ( dedicatedServer ) ;
}
}
}
Map < QueueKey , Map < ServiceInfo , List < Integer > > > serversPartitions = new HashMap < > ( ) ;
clusterRoutingService . init ( ) ;
for ( ServiceInfo ruleEngine : ruleEngines ) {
List < ServiceInfo > other = new ArrayList < > ( ruleEngines ) ;
other . removeIf ( serviceInfo - > serviceInfo . getServiceId ( ) . equals ( ruleEngine . getServiceId ( ) ) ) ;
clusterRoutingService . recalculatePartitions ( ruleEngine , other ) ;
clusterRoutingService . myPartitions . forEach ( ( queueKey , partitions ) - > {
serversPartitions . computeIfAbsent ( queueKey , k - > new HashMap < > ( ) ) . put ( ruleEngine , partitions ) ;
} ) ;
}
assertThat ( serversPartitions . keySet ( ) ) . containsAll ( queues . stream ( ) . map ( queue - > new QueueKey ( ServiceType . TB_RULE_ENGINE , queue ) ) . collect ( Collectors . toList ( ) ) ) ;
serversPartitions . forEach ( ( queueKey , partitionsPerServer ) - > {
if ( queueKey . getTenantId ( ) . isSysTenantId ( ) ) {
partitionsPerServer . forEach ( ( server , partitions ) - > {
assertThat ( server . getAssignedTenantProfilesCount ( ) ) . as ( "system queues are not assigned to dedicated servers" ) . isZero ( ) ;
} ) ;
} else {
List < ServiceInfo > responsibleServers = dedicatedServers . get ( tenants . get ( queueKey . getTenantId ( ) ) ) ;
partitionsPerServer . forEach ( ( server , partitions ) - > {
assertThat ( server . getAssignedTenantProfilesCount ( ) ) . as ( "isolated queues are only assigned to dedicated servers" ) . isPositive ( ) ;
assertThat ( responsibleServers ) . contains ( server ) ;
} ) ;
}
List < Integer > allPartitions = partitionsPerServer . values ( ) . stream ( )
. flatMap ( Collection : : stream )
. collect ( Collectors . toList ( ) ) ;
assertThat ( allPartitions ) . doesNotHaveDuplicates ( ) ;
} ) ;
}
}
ViacheslavKlimov
3 years ago