Examples with ScatterGatherRequestContext com.linkedin.pinot.transport.scattergather.ScatterGatherImpl.ScatterGatherRequestContext used on opensource projects

Example 1 with ScatterGatherRequestContext

use of com.linkedin.pinot.transport.scattergather.ScatterGatherImpl.ScatterGatherRequestContext in project pinot by linkedin.

the class ScatterGatherTest method testSelectServers.

@Test
public void testSelectServers() throws Exception {
    ExecutorService poolExecutor = MoreExecutors.sameThreadExecutor();
    ScatterGatherImpl scImpl = new ScatterGatherImpl(null, poolExecutor);
    {
        // 1 server with 2 partitions
        SegmentIdSet pg = new SegmentIdSet();
        pg.addSegment(new SegmentId("0"));
        pg.addSegment(new SegmentId("1"));
        ServerInstance serverInstance1 = new ServerInstance("localhost", 1011);
        List<ServerInstance> instances = new ArrayList<ServerInstance>();
        instances.add(serverInstance1);
        Map<ServerInstance, SegmentIdSet> pgMap = new HashMap<ServerInstance, SegmentIdSet>();
        Map<List<ServerInstance>, SegmentIdSet> invMap = new HashMap<List<ServerInstance>, SegmentIdSet>();
        pgMap.put(serverInstance1, pg);
        invMap.put(instances, pg);
        String request = "request_0";
        Map<SegmentIdSet, String> pgMapStr = new HashMap<SegmentIdSet, String>();
        pgMapStr.put(pg, request);
        ScatterGatherRequest req = new TestScatterGatherRequest(pgMap, pgMapStr);
        ScatterGatherRequestContext ctxt = new ScatterGatherRequestContext(req);
        ctxt.setInvertedMap(invMap);
        scImpl.selectServices(ctxt);
        Map<ServerInstance, SegmentIdSet> resultMap = ctxt.getSelectedServers();
        Assert.assertEquals(resultMap.size(), 1, "Count");
        Assert.assertEquals(resultMap.get(serverInstance1), pg, "Element check");
    //      System.out.println(ctxt);
    }
    {
        // 2 server with 2 partitions each
        SegmentIdSet pg = new SegmentIdSet();
        pg.addSegment(new SegmentId("0"));
        pg.addSegment(new SegmentId("1"));
        SegmentIdSet pg2 = new SegmentIdSet();
        pg2.addSegment(new SegmentId("2"));
        pg2.addSegment(new SegmentId("3"));
        ServerInstance serverInstance1 = new ServerInstance("localhost", 1011);
        ServerInstance serverInstance2 = new ServerInstance("localhost", 1012);
        List<ServerInstance> instances = new ArrayList<ServerInstance>();
        instances.add(serverInstance1);
        List<ServerInstance> instances2 = new ArrayList<ServerInstance>();
        instances2.add(serverInstance2);
        Map<ServerInstance, SegmentIdSet> pgMap = new HashMap<ServerInstance, SegmentIdSet>();
        Map<List<ServerInstance>, SegmentIdSet> invMap = new HashMap<List<ServerInstance>, SegmentIdSet>();
        pgMap.put(serverInstance1, pg);
        pgMap.put(serverInstance2, pg2);
        invMap.put(instances, pg);
        invMap.put(instances2, pg2);
        String request = "request_0";
        Map<SegmentIdSet, String> pgMapStr = new HashMap<SegmentIdSet, String>();
        pgMapStr.put(pg, request);
        ScatterGatherRequest req = new TestScatterGatherRequest(pgMap, pgMapStr);
        ScatterGatherRequestContext ctxt = new ScatterGatherRequestContext(req);
        ctxt.setInvertedMap(invMap);
        scImpl.selectServices(ctxt);
        Map<ServerInstance, SegmentIdSet> resultMap = ctxt.getSelectedServers();
        Assert.assertEquals(resultMap.size(), 2, "Count");
        Assert.assertEquals(resultMap.get(serverInstance1), pg, "Element check");
        Assert.assertEquals(resultMap.get(serverInstance2), pg2, "Element check");
    //      System.out.println(ctxt);
    }
    {
        // 2 servers sharing 2 partitions (Round-Robin selection) Partition-Group Granularity
        SegmentIdSet pg = new SegmentIdSet();
        pg.addSegment(new SegmentId("0"));
        pg.addSegment(new SegmentId("1"));
        ServerInstance serverInstance1 = new ServerInstance("localhost", 1011);
        ServerInstance serverInstance2 = new ServerInstance("localhost", 1012);
        List<ServerInstance> instances = new ArrayList<ServerInstance>();
        instances.add(serverInstance1);
        instances.add(serverInstance2);
        Map<ServerInstance, SegmentIdSet> pgMap = new HashMap<ServerInstance, SegmentIdSet>();
        Map<List<ServerInstance>, SegmentIdSet> invMap = new HashMap<List<ServerInstance>, SegmentIdSet>();
        pgMap.put(serverInstance1, pg);
        pgMap.put(serverInstance2, pg);
        invMap.put(instances, pg);
        String request = "request_0";
        Map<SegmentIdSet, String> pgMapStr = new HashMap<SegmentIdSet, String>();
        pgMapStr.put(pg, request);
        ScatterGatherRequest req = new TestScatterGatherRequest(pgMap, pgMapStr, new RoundRobinReplicaSelection(), ReplicaSelectionGranularity.SEGMENT_ID_SET, 0, 10000);
        ScatterGatherRequestContext ctxt = new ScatterGatherRequestContext(req);
        ctxt.setInvertedMap(invMap);
        scImpl.selectServices(ctxt);
        Map<ServerInstance, SegmentIdSet> resultMap = ctxt.getSelectedServers();
        Assert.assertEquals(resultMap.size(), 1, "Count");
        // first server is getting selected
        Assert.assertEquals(resultMap.get(serverInstance1), pg, "Element check");
        //      System.out.println(ctxt);
        // Run selection again. Now the second server should be selected
        scImpl.selectServices(ctxt);
        resultMap = ctxt.getSelectedServers();
        Assert.assertEquals(resultMap.size(), 1, "Count");
        // second server is getting selected
        Assert.assertEquals(resultMap.get(serverInstance2), pg, "Element check");
    //      System.out.println(ctxt);
    }
    {
        // 2 servers sharing 2 partitions (Round-Robin selection) Partition Granularity
        SegmentIdSet pg = new SegmentIdSet();
        pg.addSegment(new SegmentId("0"));
        pg.addSegment(new SegmentId("1"));
        ServerInstance serverInstance1 = new ServerInstance("localhost", 1011);
        ServerInstance serverInstance2 = new ServerInstance("localhost", 1012);
        List<ServerInstance> instances = new ArrayList<ServerInstance>();
        instances.add(serverInstance1);
        instances.add(serverInstance2);
        Map<ServerInstance, SegmentIdSet> pgMap = new HashMap<ServerInstance, SegmentIdSet>();
        Map<List<ServerInstance>, SegmentIdSet> invMap = new HashMap<List<ServerInstance>, SegmentIdSet>();
        pgMap.put(serverInstance1, pg);
        pgMap.put(serverInstance1, pg);
        invMap.put(instances, pg);
        String request = "request_0";
        Map<SegmentIdSet, String> pgMapStr = new HashMap<SegmentIdSet, String>();
        pgMapStr.put(pg, request);
        ScatterGatherRequest req = new TestScatterGatherRequest(pgMap, pgMapStr, new RoundRobinReplicaSelection(), ReplicaSelectionGranularity.SEGMENT_ID, 0, 10000);
        ScatterGatherRequestContext ctxt = new ScatterGatherRequestContext(req);
        ctxt.setInvertedMap(invMap);
        scImpl.selectServices(ctxt);
        Map<ServerInstance, SegmentIdSet> resultMap = ctxt.getSelectedServers();
        Assert.assertEquals(resultMap.size(), 2, "Count");
        // first server is getting selected
        Assert.assertFalse(resultMap.get(serverInstance1).equals(resultMap.get(serverInstance2)), "Element check");
        //      System.out.println(ctxt);
        // Run selection again. Now the second server should be selected
        scImpl.selectServices(ctxt);
        resultMap = ctxt.getSelectedServers();
        Assert.assertEquals(resultMap.size(), 2, "Count");
        // first server is getting selected
        Assert.assertFalse(resultMap.get(serverInstance1).equals(resultMap.get(serverInstance2)), "Element check");
    //      System.out.println(ctxt);
    }
}