How to use

InitializerState

in

org.keycloak.models.sessions.infinispan.initializer

private void updateLowestUnfinishedSegment() {
  this.lowestUnfinishedSegment = getNextUnfinishedSegmentFromIndex(lowestUnfinishedSegment);
}

public void markSegmentFinished(int index) {
  segments.set(index, true);
  updateLowestUnfinishedSegment();
}

while (!state.isFinished()) {
  int nodesCount = transport==null ? 1 : transport.getMembers().size();
  int distributedWorkersCount = processors * nodesCount;
  List<Integer> segments = state.getUnfinishedSegments(distributedWorkersCount);
        state.markSegmentFinished(computedSegment);
      } else {
        if (log.isTraceEnabled()) {
    log.debug("New initializer state pushed. The state is: " + state.printState());

List<SessionLoader.WorkerResult> previousResults = new LinkedList<>();
while (!state.isFinished()) {
  int nodesCount = transport==null ? 1 : transport.getMembers().size();
  int distributedWorkersCount = processors * nodesCount;
  List<Integer> segments = state.getUnfinishedSegments(distributedWorkersCount);
      state.markSegmentFinished(result.getSegment());

private InitializerState getOrCreateInitializerState() {
  InitializerState state = (InitializerState) cache.get(stateKey);
  if (state == null) {
    final int[] count = new int[1];
    // Rather use separate transactions for update and counting
    KeycloakModelUtils.runJobInTransaction(sessionFactory, new KeycloakSessionTask() {
      @Override
      public void run(KeycloakSession session) {
        sessionLoader.init(session);
      }
    });
    KeycloakModelUtils.runJobInTransaction(sessionFactory, new KeycloakSessionTask() {
      @Override
      public void run(KeycloakSession session) {
        count[0] = sessionLoader.getSessionsCount(session);
      }
    });
    state = new InitializerState();
    state.init(count[0], sessionsPerSegment);
    saveStateToCache(state);
  }
  return state;
}

  state = new InitializerState(ctx[0].getSegmentsCount());
} else {
  KeycloakModelUtils.runJobInTransaction(sessionFactory, new KeycloakSessionTask() {
    sessionLoader.toString(), ctx[0].toString(), state.toString());

private boolean isFinished() {
  InitializerState state = (InitializerState) cache.get(stateKey);
  return state != null && state.isFinished();
}

public InitializerState readObjectVersion1(ObjectInput input) throws IOException {
  return new InitializerState(
   MarshallUtil.unmarshallString(input),
   input.readInt(),
   BitSet.valueOf(MarshallUtil.unmarshallByteArray(input))
  );
}

@Override
public void writeObject(ObjectOutput output, InitializerState value) throws IOException {
  output.writeByte(VERSION_1);
  MarshallUtil.marshallString(value.getRealmId(), output);
  output.writeInt(value.segmentsCount);
  MarshallUtil.marshallByteArray(value.segments.toByteArray(), output);
}

@Override
protected boolean isFinished() {
  // Check if we should skipLoadingSessions. This can happen if someone else already did the task (For example in cross-dc environment, it was done by different DC)
  boolean isFinishedAlready = this.sessionLoader.isFinished(this);
  if (isFinishedAlready) {
    return true;
  }
  InitializerState state = getStateFromCache();
  return state != null && state.isFinished();
}

public void markSegmentFinished(int index) {
  segments.set(index);
  updateLowestUnfinishedSegment();
}

private void updateLowestUnfinishedSegment() {
  this.lowestUnfinishedSegment = getNextUnfinishedSegmentFromIndex(lowestUnfinishedSegment);
}

public InitializerState(int segmentsCount) {
  this.segmentsCount = segmentsCount;
  this.segments = new BitSet(segmentsCount);
  log.debugf("segmentsCount: %d", segmentsCount);
  updateLowestUnfinishedSegment();
}

public List<Integer> getUnfinishedSegments(int segmentCount) {
  List<Integer> result = new ArrayList<>();
  int next = lowestUnfinishedSegment;
  boolean remaining = lowestUnfinishedSegment != -1;
  while (remaining && result.size() < segmentCount) {
    next = getNextUnfinishedSegmentFromIndex(next);
    if (next == -1) {
      remaining = false;
    } else {
      result.add(next);
      next++;
    }
  }
  return result;
}

private InitializerState(String realmId, int segmentsCount, BitSet segments) {
  super(realmId);
  this.segmentsCount = segmentsCount;
  this.segments = segments;
  log.debugf("segmentsCount: %d", segmentsCount);
  updateLowestUnfinishedSegment();
}

/** Return next un-finished segments. It returns at most {@code maxSegmentCount} segments. */
public List<Integer> getUnfinishedSegments(int maxSegmentCount) {
  List<Integer> result = new LinkedList<>();
  int next = lowestUnfinishedSegment;
  boolean remaining = lowestUnfinishedSegment != -1;
  while (remaining && result.size() < maxSegmentCount) {
    next = getNextUnfinishedSegmentFromIndex(next);
    if (next == -1) {
      remaining = false;
    } else {
      result.add(next);
      next++;
    }
  }
  return result;
}

public void init(int sessionsCount, int sessionsPerSegment) {
  this.sessionsCount = sessionsCount;
  int segmentsCount = sessionsCount / sessionsPerSegment;
  if (sessionsPerSegment * segmentsCount < sessionsCount) {
    segmentsCount = segmentsCount + 1;
  }
  log.debugf("sessionsCount: %d, sessionsPerSegment: %d, segmentsCount: %d", sessionsCount, sessionsPerSegment, segmentsCount);
  for (int i=0 ; i<segmentsCount ; i++) {
    segments.add(false);
  }
  updateLowestUnfinishedSegment();
}