mirror of https://github.com/apache/druid.git
allow server selection to be aware of query (#10428)
* add query through to server selector * add nullable extensions, deprecate old methods with defaults * style changes * add nullable to ServerSelectorStrategy * fix test coverage * missing override in test * add null check
This commit is contained in:
parent
6f2ce8f0a5
commit
ca3b925133
|
@ -588,7 +588,7 @@ public class CachingClusteredClient implements QuerySegmentWalker
|
|||
{
|
||||
final SortedMap<DruidServer, List<SegmentDescriptor>> serverSegments = new TreeMap<>();
|
||||
for (SegmentServerSelector segmentServer : segments) {
|
||||
final QueryableDruidServer queryableDruidServer = segmentServer.getServer().pick();
|
||||
final QueryableDruidServer queryableDruidServer = segmentServer.getServer().pick(query);
|
||||
|
||||
if (queryableDruidServer == null) {
|
||||
log.makeAlert(
|
||||
|
@ -812,7 +812,8 @@ public class CachingClusteredClient implements QuerySegmentWalker
|
|||
Hasher hasher = Hashing.sha1().newHasher();
|
||||
boolean hasOnlyHistoricalSegments = true;
|
||||
for (SegmentServerSelector p : segments) {
|
||||
if (!p.getServer().pick().getServer().isSegmentReplicationTarget()) {
|
||||
QueryableDruidServer queryableServer = p.getServer().pick(query);
|
||||
if (queryableServer == null || !queryableServer.getServer().isSegmentReplicationTarget()) {
|
||||
hasOnlyHistoricalSegments = false;
|
||||
break;
|
||||
}
|
||||
|
|
|
@ -21,6 +21,7 @@ package org.apache.druid.client.selector;
|
|||
|
||||
import com.google.common.collect.Iterables;
|
||||
import it.unimi.dsi.fastutil.ints.Int2ObjectRBTreeMap;
|
||||
import org.apache.druid.query.Query;
|
||||
import org.apache.druid.timeline.DataSegment;
|
||||
|
||||
import javax.annotation.Nullable;
|
||||
|
@ -41,16 +42,18 @@ public abstract class AbstractTierSelectorStrategy implements TierSelectorStrate
|
|||
|
||||
@Nullable
|
||||
@Override
|
||||
public QueryableDruidServer pick(
|
||||
public <T> QueryableDruidServer pick(
|
||||
Query<T> query,
|
||||
Int2ObjectRBTreeMap<Set<QueryableDruidServer>> prioritizedServers,
|
||||
DataSegment segment
|
||||
)
|
||||
{
|
||||
return Iterables.getOnlyElement(pick(prioritizedServers, segment, 1), null);
|
||||
return Iterables.getOnlyElement(pick(query, prioritizedServers, segment, 1), null);
|
||||
}
|
||||
|
||||
@Override
|
||||
public List<QueryableDruidServer> pick(
|
||||
public <T> List<QueryableDruidServer> pick(
|
||||
Query<T> query,
|
||||
Int2ObjectRBTreeMap<Set<QueryableDruidServer>> prioritizedServers,
|
||||
DataSegment segment,
|
||||
int numServersToPick
|
||||
|
@ -58,7 +61,7 @@ public abstract class AbstractTierSelectorStrategy implements TierSelectorStrate
|
|||
{
|
||||
List<QueryableDruidServer> result = new ArrayList<>(numServersToPick);
|
||||
for (Set<QueryableDruidServer> priorityServers : prioritizedServers.values()) {
|
||||
result.addAll(serverSelectorStrategy.pick(priorityServers, segment, numServersToPick - result.size()));
|
||||
result.addAll(serverSelectorStrategy.pick(query, priorityServers, segment, numServersToPick - result.size()));
|
||||
if (result.size() == numServersToPick) {
|
||||
break;
|
||||
}
|
||||
|
|
|
@ -24,6 +24,7 @@ import com.google.common.collect.Ordering;
|
|||
import org.apache.druid.client.DirectDruidClient;
|
||||
import org.apache.druid.timeline.DataSegment;
|
||||
|
||||
import javax.annotation.Nullable;
|
||||
import java.util.Collections;
|
||||
import java.util.Comparator;
|
||||
import java.util.List;
|
||||
|
@ -34,6 +35,7 @@ public class ConnectionCountServerSelectorStrategy implements ServerSelectorStra
|
|||
private static final Comparator<QueryableDruidServer> COMPARATOR =
|
||||
Comparator.comparingInt(s -> ((DirectDruidClient) s.getQueryRunner()).getNumOpenConnections());
|
||||
|
||||
@Nullable
|
||||
@Override
|
||||
public QueryableDruidServer pick(Set<QueryableDruidServer> servers, DataSegment segment)
|
||||
{
|
||||
|
|
|
@ -24,6 +24,7 @@ import com.google.common.collect.Iterators;
|
|||
import com.google.common.collect.Lists;
|
||||
import org.apache.druid.timeline.DataSegment;
|
||||
|
||||
import javax.annotation.Nullable;
|
||||
import java.util.Collections;
|
||||
import java.util.List;
|
||||
import java.util.Set;
|
||||
|
@ -31,6 +32,7 @@ import java.util.concurrent.ThreadLocalRandom;
|
|||
|
||||
public class RandomServerSelectorStrategy implements ServerSelectorStrategy
|
||||
{
|
||||
@Nullable
|
||||
@Override
|
||||
public QueryableDruidServer pick(Set<QueryableDruidServer> servers, DataSegment segment)
|
||||
{
|
||||
|
|
|
@ -21,6 +21,7 @@ package org.apache.druid.client.selector;
|
|||
|
||||
import it.unimi.dsi.fastutil.ints.Int2ObjectRBTreeMap;
|
||||
import org.apache.druid.client.DataSegmentInterner;
|
||||
import org.apache.druid.query.Query;
|
||||
import org.apache.druid.server.coordination.DruidServerMetadata;
|
||||
import org.apache.druid.server.coordination.ServerType;
|
||||
import org.apache.druid.timeline.DataSegment;
|
||||
|
@ -36,7 +37,7 @@ import java.util.concurrent.atomic.AtomicReference;
|
|||
|
||||
/**
|
||||
*/
|
||||
public class ServerSelector implements DiscoverySelector<QueryableDruidServer>, Overshadowable<ServerSelector>
|
||||
public class ServerSelector implements Overshadowable<ServerSelector>
|
||||
{
|
||||
|
||||
private final Int2ObjectRBTreeMap<Set<QueryableDruidServer>> historicalServers;
|
||||
|
@ -160,14 +161,13 @@ public class ServerSelector implements DiscoverySelector<QueryableDruidServer>,
|
|||
}
|
||||
|
||||
@Nullable
|
||||
@Override
|
||||
public QueryableDruidServer pick()
|
||||
public <T> QueryableDruidServer pick(@Nullable Query<T> query)
|
||||
{
|
||||
synchronized (this) {
|
||||
if (!historicalServers.isEmpty()) {
|
||||
return strategy.pick(historicalServers, segment.get());
|
||||
return strategy.pick(query, historicalServers, segment.get());
|
||||
}
|
||||
return strategy.pick(realtimeServers, segment.get());
|
||||
return strategy.pick(query, realtimeServers, segment.get());
|
||||
}
|
||||
}
|
||||
|
||||
|
|
|
@ -21,8 +21,11 @@ package org.apache.druid.client.selector;
|
|||
|
||||
import com.fasterxml.jackson.annotation.JsonSubTypes;
|
||||
import com.fasterxml.jackson.annotation.JsonTypeInfo;
|
||||
import com.google.common.collect.Iterables;
|
||||
import org.apache.druid.query.Query;
|
||||
import org.apache.druid.timeline.DataSegment;
|
||||
|
||||
import javax.annotation.Nullable;
|
||||
import java.util.List;
|
||||
import java.util.Set;
|
||||
|
||||
|
@ -33,7 +36,29 @@ import java.util.Set;
|
|||
})
|
||||
public interface ServerSelectorStrategy
|
||||
{
|
||||
QueryableDruidServer pick(Set<QueryableDruidServer> servers, DataSegment segment);
|
||||
@Nullable
|
||||
default <T> QueryableDruidServer pick(@Nullable Query<T> query, Set<QueryableDruidServer> servers, DataSegment segment)
|
||||
{
|
||||
return Iterables.getOnlyElement(pick(query, servers, segment, 1), null);
|
||||
}
|
||||
|
||||
default <T> List<QueryableDruidServer> pick(@Nullable Query<T> query, Set<QueryableDruidServer> servers, DataSegment segment,
|
||||
int numServersToPick)
|
||||
{
|
||||
return pick(servers, segment, numServersToPick);
|
||||
}
|
||||
|
||||
@Deprecated
|
||||
@Nullable
|
||||
default QueryableDruidServer pick(Set<QueryableDruidServer> servers, DataSegment segment)
|
||||
{
|
||||
return pick(null, servers, segment);
|
||||
}
|
||||
|
||||
@Deprecated
|
||||
default List<QueryableDruidServer> pick(Set<QueryableDruidServer> servers, DataSegment segment, int numServersToPick)
|
||||
{
|
||||
return pick(null, servers, segment, numServersToPick);
|
||||
}
|
||||
|
||||
List<QueryableDruidServer> pick(Set<QueryableDruidServer> servers, DataSegment segment, int numServersToPick);
|
||||
}
|
||||
|
|
|
@ -22,9 +22,11 @@ package org.apache.druid.client.selector;
|
|||
import com.fasterxml.jackson.annotation.JsonSubTypes;
|
||||
import com.fasterxml.jackson.annotation.JsonTypeInfo;
|
||||
import it.unimi.dsi.fastutil.ints.Int2ObjectRBTreeMap;
|
||||
import org.apache.druid.query.Query;
|
||||
import org.apache.druid.timeline.DataSegment;
|
||||
|
||||
import javax.annotation.Nullable;
|
||||
|
||||
import java.util.Comparator;
|
||||
import java.util.List;
|
||||
import java.util.Set;
|
||||
|
@ -41,12 +43,36 @@ public interface TierSelectorStrategy
|
|||
{
|
||||
Comparator<Integer> getComparator();
|
||||
|
||||
@Deprecated
|
||||
@Nullable
|
||||
QueryableDruidServer pick(Int2ObjectRBTreeMap<Set<QueryableDruidServer>> prioritizedServers, DataSegment segment);
|
||||
default QueryableDruidServer pick(Int2ObjectRBTreeMap<Set<QueryableDruidServer>> prioritizedServers, DataSegment segment)
|
||||
{
|
||||
return pick(null, prioritizedServers, segment);
|
||||
}
|
||||
|
||||
List<QueryableDruidServer> pick(
|
||||
@Deprecated
|
||||
default List<QueryableDruidServer> pick(
|
||||
Int2ObjectRBTreeMap<Set<QueryableDruidServer>> prioritizedServers,
|
||||
DataSegment segment,
|
||||
int numServersToPick
|
||||
);
|
||||
int numServersToPick)
|
||||
{
|
||||
return pick(null, prioritizedServers, segment, numServersToPick);
|
||||
}
|
||||
|
||||
@Nullable
|
||||
default <T> QueryableDruidServer pick(@Nullable Query<T> query,
|
||||
Int2ObjectRBTreeMap<Set<QueryableDruidServer>> prioritizedServers,
|
||||
DataSegment segment)
|
||||
{
|
||||
return pick(prioritizedServers, segment);
|
||||
}
|
||||
|
||||
default <T> List<QueryableDruidServer> pick(
|
||||
@Nullable Query<T> query,
|
||||
Int2ObjectRBTreeMap<Set<QueryableDruidServer>> prioritizedServers,
|
||||
DataSegment segment,
|
||||
int numServersToPick)
|
||||
{
|
||||
return pick(prioritizedServers, segment, numServersToPick);
|
||||
}
|
||||
}
|
||||
|
|
|
@ -135,7 +135,7 @@ public class BrokerServerViewTest extends CuratorTestBase
|
|||
ServerSelector selector = (actualPartitionHolder.iterator().next()).getObject();
|
||||
Assert.assertFalse(selector.isEmpty());
|
||||
Assert.assertEquals(segment, selector.getSegment());
|
||||
Assert.assertEquals(druidServer, selector.pick().getServer());
|
||||
Assert.assertEquals(druidServer, selector.pick(null).getServer());
|
||||
|
||||
unannounceSegmentForServer(druidServer, segment, zkPathsConfig);
|
||||
Assert.assertTrue(timing.forWaiting().awaitLatch(segmentRemovedLatch));
|
||||
|
@ -384,7 +384,7 @@ public class BrokerServerViewTest extends CuratorTestBase
|
|||
ServerSelector selector = ((SingleElementPartitionChunk<ServerSelector>) actualPartitionHolder.iterator()
|
||||
.next()).getObject();
|
||||
Assert.assertFalse(selector.isEmpty());
|
||||
Assert.assertEquals(expectedPair.rhs.rhs.lhs, selector.pick().getServer());
|
||||
Assert.assertEquals(expectedPair.rhs.rhs.lhs, selector.pick(null).getServer());
|
||||
Assert.assertEquals(expectedPair.rhs.rhs.rhs, selector.getSegment());
|
||||
}
|
||||
}
|
||||
|
|
|
@ -332,7 +332,7 @@ public class CachingClusteredClientCacheKeyManagerTest extends EasyMockSupport
|
|||
0
|
||||
);
|
||||
expect(server.isSegmentReplicationTarget()).andReturn(isHistorical).anyTimes();
|
||||
expect(serverSelector.pick()).andReturn(queryableDruidServer).anyTimes();
|
||||
expect(serverSelector.pick(query)).andReturn(queryableDruidServer).anyTimes();
|
||||
expect(queryableDruidServer.getServer()).andReturn(server).anyTimes();
|
||||
expect(serverSelector.getSegment()).andReturn(segment).anyTimes();
|
||||
replay(serverSelector, queryableDruidServer, server);
|
||||
|
|
|
@ -223,7 +223,7 @@ public class DirectDruidClientTest
|
|||
|
||||
Assert.assertEquals(2, client2.getNumOpenConnections());
|
||||
|
||||
Assert.assertEquals(serverSelector.pick(), queryableDruidServer2);
|
||||
Assert.assertEquals(serverSelector.pick(null), queryableDruidServer2);
|
||||
|
||||
EasyMock.verify(httpClient);
|
||||
}
|
||||
|
|
|
@ -23,6 +23,7 @@ import org.apache.druid.client.DirectDruidClient;
|
|||
import org.apache.druid.client.DruidServer;
|
||||
import org.apache.druid.java.util.common.DateTimes;
|
||||
import org.apache.druid.java.util.common.Intervals;
|
||||
import org.apache.druid.query.Query;
|
||||
import org.apache.druid.server.coordination.DruidServerMetadata;
|
||||
import org.apache.druid.server.coordination.ServerType;
|
||||
import org.apache.druid.timeline.DataSegment;
|
||||
|
@ -31,15 +32,37 @@ import org.easymock.EasyMock;
|
|||
import org.junit.Assert;
|
||||
import org.junit.Test;
|
||||
|
||||
import javax.annotation.Nullable;
|
||||
import java.util.ArrayList;
|
||||
import java.util.Arrays;
|
||||
import java.util.Collections;
|
||||
import java.util.HashMap;
|
||||
import java.util.HashSet;
|
||||
import java.util.List;
|
||||
import java.util.Set;
|
||||
|
||||
public class TierSelectorStrategyTest
|
||||
{
|
||||
|
||||
@Test
|
||||
public void testHighestPriorityTierSelectorStrategyRealtime()
|
||||
{
|
||||
DirectDruidClient client = EasyMock.createMock(DirectDruidClient.class);
|
||||
QueryableDruidServer lowPriority = new QueryableDruidServer(
|
||||
new DruidServer("test1", "localhost", null, 0, ServerType.REALTIME, DruidServer.DEFAULT_TIER, 0),
|
||||
client
|
||||
);
|
||||
QueryableDruidServer highPriority = new QueryableDruidServer(
|
||||
new DruidServer("test1", "localhost", null, 0, ServerType.REALTIME, DruidServer.DEFAULT_TIER, 1),
|
||||
client
|
||||
);
|
||||
|
||||
testTierSelectorStrategy(
|
||||
new HighestPriorityTierSelectorStrategy(new ConnectionCountServerSelectorStrategy()),
|
||||
highPriority, lowPriority
|
||||
);
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testHighestPriorityTierSelectorStrategy()
|
||||
{
|
||||
|
@ -168,9 +191,7 @@ public class TierSelectorStrategyTest
|
|||
new DruidServer("test1", "localhost", null, 0, ServerType.HISTORICAL, DruidServer.DEFAULT_TIER, 3),
|
||||
client
|
||||
);
|
||||
|
||||
testTierSelectorStrategy(
|
||||
new CustomTierSelectorStrategy(
|
||||
TierSelectorStrategy tierSelectorStrategy = new CustomTierSelectorStrategy(
|
||||
new ConnectionCountServerSelectorStrategy(),
|
||||
new CustomTierSelectorStrategyConfig()
|
||||
{
|
||||
|
@ -180,7 +201,9 @@ public class TierSelectorStrategyTest
|
|||
return Arrays.asList(2, 0, -1);
|
||||
}
|
||||
}
|
||||
),
|
||||
);
|
||||
testTierSelectorStrategy(
|
||||
tierSelectorStrategy,
|
||||
p3, p1, p0, p4, p2
|
||||
);
|
||||
}
|
||||
|
@ -216,9 +239,35 @@ public class TierSelectorStrategyTest
|
|||
serverSelector.addServerAndUpdateSegment(server, serverSelector.getSegment());
|
||||
}
|
||||
|
||||
Assert.assertEquals(expectedSelection[0], serverSelector.pick());
|
||||
Assert.assertEquals(expectedSelection[0], serverSelector.pick(null));
|
||||
Assert.assertEquals(expectedSelection[0], serverSelector.pick(EasyMock.createMock(Query.class)));
|
||||
Assert.assertEquals(expectedCandidates, serverSelector.getCandidates(-1));
|
||||
Assert.assertEquals(expectedCandidates.subList(0, 2), serverSelector.getCandidates(2));
|
||||
}
|
||||
|
||||
@Test
|
||||
public void testServerSelectorStrategyDefaults()
|
||||
{
|
||||
DirectDruidClient client = EasyMock.createMock(DirectDruidClient.class);
|
||||
QueryableDruidServer p0 = new QueryableDruidServer(
|
||||
new DruidServer("test1", "localhost", null, 0, ServerType.HISTORICAL, DruidServer.DEFAULT_TIER, -1),
|
||||
client
|
||||
);
|
||||
Set<QueryableDruidServer> servers = new HashSet<>();
|
||||
servers.add(p0);
|
||||
RandomServerSelectorStrategy strategy = new RandomServerSelectorStrategy();
|
||||
Assert.assertEquals(strategy.pick(servers, EasyMock.createMock(DataSegment.class)), p0);
|
||||
Assert.assertEquals(strategy.pick(EasyMock.createMock(Query.class), servers, EasyMock.createMock(DataSegment.class)), p0);
|
||||
ServerSelectorStrategy defaultDeprecatedServerSelectorStrategy = new ServerSelectorStrategy() {
|
||||
@Override
|
||||
public <T> List<QueryableDruidServer> pick(@Nullable Query<T> query, Set<QueryableDruidServer> servers, DataSegment segment,
|
||||
int numServersToPick)
|
||||
{
|
||||
return strategy.pick(servers, segment, numServersToPick);
|
||||
}
|
||||
};
|
||||
Assert.assertEquals(defaultDeprecatedServerSelectorStrategy.pick(servers, EasyMock.createMock(DataSegment.class)), p0);
|
||||
Assert.assertEquals(defaultDeprecatedServerSelectorStrategy.pick(servers, EasyMock.createMock(DataSegment.class), 1).get(0), p0);
|
||||
}
|
||||
|
||||
}
|
||||
|
|
Loading…
Reference in New Issue