Code review comment: Timestamp is the sort key here. As our OrderedListState interface is based on TimestampedValue<T>, the sort key is actually an Instant, but I agree that I should be consistent here by using "sort key".  The reason why I put the value encoded bytes before the sort key encoded bytes is that the same order is used in the coder of TimestampedValue<T>. https://github.com/apache/beam/blob/0d46e304f176847e897eef244c50cb29af8c6451/sdks/java/core/src/main/java/org/apache/beam/sdk/values/TimestampedValue.java#L110  Here, we pack the data into the output stream as if we are calling `TimestampedValue.encode()`. However, we cannot call encode() directly, because we don't want to decode the value and encode it again here.  (In my previous attempt, I did try to decode the value and encode it again. It seems to introduce unnecessary overhead. https://github.com/apache/beam/pull/30317/commits/d06e0a5ebe8405334112a2d797bb5c9571b46b9c)

Code before review:
public CompletableFuture<StateResponse> handle(StateRequest.Builder requestBuilder) {
assertEquals("", requestBuilder.getId());
requestBuilder.setId(generateId());
StateRequest request = requestBuilder.build();
StateKey key = request.getStateKey();
StateResponse.Builder response;
assertNotEquals(RequestCase.REQUEST_NOT_SET, request.getRequestCase());
assertNotEquals(TypeCase.TYPE_NOT_SET, key.getTypeCase());
if (key.getTypeCase() == TypeCase.MULTIMAP_SIDE_INPUT || key.getTypeCase() == TypeCase.RUNNER) {
assertEquals(GET, request.getRequestCase());
}
if (key.getTypeCase() == TypeCase.MULTIMAP_KEYS_VALUES_SIDE_INPUT && !data.containsKey(key)) {
throw new UnsupportedOperationException("No multimap keys values states provided.");
}
switch (request.getRequestCase()) {
case GET:
{
List<ByteString> byteStrings =
data.getOrDefault(request.getStateKey(), Collections.singletonList(ByteString.EMPTY));
int block = 0;
if (request.getGet().getContinuationToken().size() > 0) {
block = Integer.parseInt(request.getGet().getContinuationToken().toStringUtf8());
}
ByteString returnBlock = byteStrings.get(block);
ByteString continuationToken = ByteString.EMPTY;
if (byteStrings.size() > block + 1) {
continuationToken = ByteString.copyFromUtf8(Integer.toString(block + 1));
}
response =
StateResponse.newBuilder()
.setGet(
StateGetResponse.newBuilder()
.setData(returnBlock)
.setContinuationToken(continuationToken));
}
break;
case CLEAR:
data.remove(request.getStateKey());
response = StateResponse.newBuilder().setClear(StateClearResponse.getDefaultInstance());
break;
case APPEND:
List<ByteString> previousValue =
data.computeIfAbsent(request.getStateKey(), (unused) -> new ArrayList<>());
previousValue.add(request.getAppend().getData());
response = StateResponse.newBuilder().setAppend(StateAppendResponse.getDefaultInstance());
break;
case ORDERED_LIST_GET:
{
long start = request.getOrderedListGet().getRange().getStart();
long end = request.getOrderedListGet().getRange().getEnd();
KvCoder<Long, Integer> coder = KvCoder.of(VarLongCoder.of(), VarIntCoder.of());
long sortKey = start;
int index = 0;
if (request.getOrderedListGet().getContinuationToken().size() > 0) {
try {
KV<Long, Integer> cursor =
coder.decode(request.getOrderedListGet().getContinuationToken().newInput());
sortKey = cursor.getKey();
index = cursor.getValue();
} catch (IOException e) {
throw new RuntimeException(e);
}
}
ByteString continuationToken;
ByteString returnBlock = ByteString.EMPTY;
;
try {
if (sortKey < start || sortKey >= end) {
throw new IndexOutOfBoundsException("sort key out of range");
}
NavigableSet<Long> subset =
orderedListKeys
.getOrDefault(request.getStateKey(), new TreeSet<>())
.subSet(sortKey, true, end, false);
Long nextSortKey = subset.first();
StateKey.Builder keyBuilder = request.getStateKey().toBuilder();
keyBuilder.getOrderedListUserStateBuilder().setSortKey(nextSortKey);
List<ByteString> byteStrings =
data.getOrDefault(keyBuilder.build(), Collections.singletonList(ByteString.EMPTY));
returnBlock = byteStrings.get(index);
if (byteStrings.size() > index + 1) {
index += 1;
} else {
nextSortKey = subset.tailSet(nextSortKey, false).first();
index = 0;
}
ByteStringOutputStream outputStream = new ByteStringOutputStream();
try {
KV<Long, Integer> cursor = KV.of(nextSortKey, index);
coder.encode(cursor, outputStream);
} catch (IOException e) {
throw new RuntimeException(e);
}
continuationToken = outputStream.toByteString();
} catch (NoSuchElementException | IndexOutOfBoundsException e) {
continuationToken = ByteString.EMPTY;
}
response =
StateResponse.newBuilder()
.setOrderedListGet(
OrderedListStateGetResponse.newBuilder()
.setData(returnBlock)
.setContinuationToken(continuationToken));
}
break;
case ORDERED_LIST_UPDATE:
for (OrderedListRange r : request.getOrderedListUpdate().getDeletesList()) {
List<Long> keysToRemove =
new ArrayList<>(
orderedListKeys
.getOrDefault(request.getStateKey(), new TreeSet<>())
.subSet(r.getStart(), true, r.getEnd(), false));
for (Long l : keysToRemove) {
StateKey.Builder keyBuilder = request.getStateKey().toBuilder();
keyBuilder.getOrderedListUserStateBuilder().setSortKey(l);
data.remove(keyBuilder.build());
orderedListKeys.get(request.getStateKey()).remove(l);
}
}
for (OrderedListEntry e : request.getOrderedListUpdate().getInsertsList()) {
StateKey.Builder keyBuilder = request.getStateKey().toBuilder();
keyBuilder.getOrderedListUserStateBuilder().setSortKey(e.getSortKey());
ByteStringOutputStream outStream = new ByteStringOutputStream();
try {
InstantCoder.of().encode(Instant.ofEpochMilli(e.getSortKey()), outStream);
} catch (IOException ex) {
throw new RuntimeException(ex);
}
ByteString output = e.getData().concat(outStream.toByteString());
List<ByteString> previousValues =
data.computeIfAbsent(keyBuilder.build(), (unused) -> new ArrayList<>());
previousValues.add(output);
orderedListKeys
.computeIfAbsent(request.getStateKey(), (unused) -> new TreeSet<>())
.add(e.getSortKey());
}
response =
StateResponse.newBuilder()
.setOrderedListUpdate(OrderedListStateUpdateResponse.getDefaultInstance());
break;
default:
throw new IllegalStateException(
String.format("Unknown request type %s", request.getRequestCase()));
}
return CompletableFuture.completedFuture(response.setId(requestBuilder.getId()).build());
}


Code after addressing the review:
public CompletableFuture<StateResponse> handle(StateRequest.Builder requestBuilder) {
assertEquals("", requestBuilder.getId());
requestBuilder.setId(generateId());
StateRequest request = requestBuilder.build();
StateKey key = request.getStateKey();
StateResponse.Builder response;
assertNotEquals(RequestCase.REQUEST_NOT_SET, request.getRequestCase());
assertNotEquals(TypeCase.TYPE_NOT_SET, key.getTypeCase());
if (key.getTypeCase() == TypeCase.MULTIMAP_SIDE_INPUT || key.getTypeCase() == TypeCase.RUNNER) {
assertEquals(GET, request.getRequestCase());
}
if (key.getTypeCase() == TypeCase.MULTIMAP_KEYS_VALUES_SIDE_INPUT && !data.containsKey(key)) {
throw new UnsupportedOperationException("No multimap keys values states provided.");
}
switch (request.getRequestCase()) {
case GET:
if (key.getTypeCase() == TypeCase.ORDERED_LIST_USER_STATE) {
long start = key.getOrderedListUserState().getRange().getStart();
long end = key.getOrderedListUserState().getRange().getEnd();
KvCoder<Long, Integer> coder = KvCoder.of(VarLongCoder.of(), VarIntCoder.of());
long sortKey = start;
int index = 0;
if (!request.getGet().getContinuationToken().isEmpty()) {
try {
KV<Long, Integer> cursor =
coder.decode(request.getGet().getContinuationToken().newInput());
sortKey = cursor.getKey();
index = cursor.getValue();
} catch (IOException e) {
throw new RuntimeException(e);
}
}
ByteString continuationToken;
ByteString returnBlock = ByteString.EMPTY;
try {
if (sortKey < start || sortKey >= end) {
throw new IndexOutOfBoundsException("sort key out of range");
}
StateKey.Builder stateKeyWithoutRange = request.getStateKey().toBuilder();
stateKeyWithoutRange.getOrderedListUserStateBuilder().clearRange();
NavigableSet<Long> subset =
orderedListSortKeysFromStateKey
.getOrDefault(stateKeyWithoutRange.build(), new TreeSet<>())
.subSet(sortKey, true, end, false);
Long nextSortKey = subset.first();
StateKey.Builder keyBuilder = request.getStateKey().toBuilder();
keyBuilder
.getOrderedListUserStateBuilder()
.getRangeBuilder()
.setStart(nextSortKey)
.setEnd(nextSortKey + 1);
List<ByteString> byteStrings =
data.getOrDefault(keyBuilder.build(), Collections.singletonList(ByteString.EMPTY));
returnBlock = byteStrings.get(index);
if (byteStrings.size() > index + 1) {
index += 1;
} else {
nextSortKey = subset.tailSet(nextSortKey, false).first();
index = 0;
}
ByteStringOutputStream outputStream = new ByteStringOutputStream();
try {
KV<Long, Integer> cursor = KV.of(nextSortKey, index);
coder.encode(cursor, outputStream);
} catch (IOException e) {
throw new RuntimeException(e);
}
continuationToken = outputStream.toByteString();
} catch (NoSuchElementException | IndexOutOfBoundsException e) {
continuationToken = ByteString.EMPTY;
}
response =
StateResponse.newBuilder()
.setGet(
StateGetResponse.newBuilder()
.setData(returnBlock)
.setContinuationToken(continuationToken));
} else {
List<ByteString> byteStrings =
data.getOrDefault(request.getStateKey(), Collections.singletonList(ByteString.EMPTY));
int block = 0;
if (!request.getGet().getContinuationToken().isEmpty()) {
block = Integer.parseInt(request.getGet().getContinuationToken().toStringUtf8());
}
ByteString returnBlock = byteStrings.get(block);
ByteString continuationToken = ByteString.EMPTY;
if (byteStrings.size() > block + 1) {
continuationToken = ByteString.copyFromUtf8(Integer.toString(block + 1));
}
response =
StateResponse.newBuilder()
.setGet(
StateGetResponse.newBuilder()
.setData(returnBlock)
.setContinuationToken(continuationToken));
}
break;
case CLEAR:
if (key.getTypeCase() == TypeCase.ORDERED_LIST_USER_STATE) {
OrderedListRange r = request.getStateKey().getOrderedListUserState().getRange();
StateKey.Builder stateKeyWithoutRange = request.getStateKey().toBuilder();
stateKeyWithoutRange.getOrderedListUserStateBuilder().clearRange();
List<Long> keysToRemove =
new ArrayList<>(
orderedListSortKeysFromStateKey
.getOrDefault(stateKeyWithoutRange.build(), new TreeSet<>())
.subSet(r.getStart(), true, r.getEnd(), false));
for (Long l : keysToRemove) {
StateKey.Builder keyBuilder = request.getStateKey().toBuilder();
keyBuilder.getOrderedListUserStateBuilder().getRangeBuilder().setStart(l).setEnd(l + 1);
data.remove(keyBuilder.build());
orderedListSortKeysFromStateKey.get(stateKeyWithoutRange.build()).remove(l);
}
} else {
data.remove(request.getStateKey());
}
response = StateResponse.newBuilder().setClear(StateClearResponse.getDefaultInstance());
break;
case APPEND:
if (key.getTypeCase() == TypeCase.ORDERED_LIST_USER_STATE) {
InputStream inStream = request.getAppend().getData().newInput();
TimestampedValueCoder<byte[]> coder = TimestampedValueCoder.of(ByteArrayCoder.of());
try {
while (inStream.available() > 0) {
TimestampedValue<byte[]> tv = coder.decode(inStream);
ByteStringOutputStream outStream = new ByteStringOutputStream();
coder.encode(tv, outStream);
ByteString output = outStream.toByteString();
StateKey.Builder keyBuilder = request.getStateKey().toBuilder();
long sortKey = tv.getTimestamp().getMillis();
keyBuilder
.getOrderedListUserStateBuilder()
.getRangeBuilder()
.setStart(sortKey)
.setEnd(sortKey + 1);
List<ByteString> previousValues =
data.computeIfAbsent(keyBuilder.build(), (unused) -> new ArrayList<>());
previousValues.add(output);
StateKey.Builder stateKeyWithoutRange = request.getStateKey().toBuilder();
stateKeyWithoutRange.getOrderedListUserStateBuilder().clearRange();
orderedListSortKeysFromStateKey
.computeIfAbsent(stateKeyWithoutRange.build(), (unused) -> new TreeSet<>())
.add(sortKey);
}
} catch (IOException ex) {
throw new RuntimeException(ex);
}
} else {
List<ByteString> previousValue =
data.computeIfAbsent(request.getStateKey(), (unused) -> new ArrayList<>());
previousValue.add(request.getAppend().getData());
}
response = StateResponse.newBuilder().setAppend(StateAppendResponse.getDefaultInstance());
break;
default:
throw new IllegalStateException(
String.format("Unknown request type %s", request.getRequestCase()));
}
return CompletableFuture.completedFuture(response.setId(requestBuilder.getId()).build());
}