CostBalancerStrategy.pickSegmentToMove can pick the same segment twice

[egor-ryashin]

  @Override
  public BalancerSegmentHolder pickSegmentToMove(final List<ServerHolder> serverHolders)
  {
    ReservoirSegmentSampler sampler = new ReservoirSegmentSampler();
    return sampler.getRandomBalancerSegmentHolder(serverHolders);
  }

I see a probability here that the coordinator picks the same segment several times within a single rebalancing phase as nothing keeps the record of selected segments.

[leventov]

It doesn't seem to me that this is a "bug" in the sense that the authors of the original code didn't realize this thing. Probably it was considered "good enough". But it could be improved.

[egor-ryashin]

BTW, additional tests show that pickSegmentToMove efficiency substantially decreases only for small amount of segments (20% duplicates for 500 segments) and negligible (<10 if I remember correctly) for 100000 segments.
So it's mostly confusing during tests.

[egor-ryashin]

BTW, I believe, it can pick segments and request to load them while they are requested for drop during a LoadRule run.

[kfaraz]

This is now fixed because the implementation takes the reservoir size and returns an iterator over those segments in a one-shot sampling.

default Iterator<BalancerSegmentHolder> pickSegmentsToMove(
      List<ServerHolder> serverHolders,
      Set<String> broadcastDatasources,
      int reservoirSize
  );

static List<BalancerSegmentHolder> getRandomBalancerSegmentHolders(
      final List<ServerHolder> serverHolders,
      Set<String> broadcastDatasources,
      int k
  )
  {
    List<BalancerSegmentHolder> holders = new ArrayList<>(k);
    int numSoFar = 0;

    for (ServerHolder server : serverHolders) {
      if (!server.getServer().getType().isSegmentReplicationTarget()) {
        // if the server only handles broadcast segments (which don't need to be rebalanced), we have nothing to do
        continue;
      }

      for (DataSegment segment : server.getServer().iterateAllSegments()) {
        if (broadcastDatasources.contains(segment.getDataSource())) {
          // we don't need to rebalance segments that were assigned via broadcast rules
          continue;
        }

        if (numSoFar < k) {
          holders.add(new BalancerSegmentHolder(server.getServer(), segment));
          numSoFar++;
          continue;
        }
        int randNum = ThreadLocalRandom.current().nextInt(numSoFar + 1);
        if (randNum < k) {
          holders.set(randNum, new BalancerSegmentHolder(server.getServer(), segment));
        }
        numSoFar++;
      }
    }
    return holders;
  }