diff --git a/atlasdb-impl-shared/src/main/java/com/palantir/atlasdb/sweep/asts/bucketingthings/BucketAssignerState.java b/atlasdb-impl-shared/src/main/java/com/palantir/atlasdb/sweep/asts/bucketingthings/BucketAssignerState.java new file mode 100644 index 0000000000..62f2f6ed93 --- /dev/null +++ b/atlasdb-impl-shared/src/main/java/com/palantir/atlasdb/sweep/asts/bucketingthings/BucketAssignerState.java @@ -0,0 +1,152 @@ +/* + * (c) Copyright 2024 Palantir Technologies Inc. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.palantir.atlasdb.sweep.asts.bucketingthings; + +import com.fasterxml.jackson.annotation.JsonSubTypes; +import com.fasterxml.jackson.annotation.JsonTypeInfo; +import com.fasterxml.jackson.databind.annotation.JsonDeserialize; +import com.fasterxml.jackson.databind.annotation.JsonSerialize; +import org.immutables.value.Value; + +@JsonTypeInfo(use = JsonTypeInfo.Id.NAME, property = "type") +@JsonSubTypes({ + @JsonSubTypes.Type(value = ImmutableStart.class, name = BucketAssignerState.Start.TYPE), + @JsonSubTypes.Type(value = ImmutableOpening.class, name = BucketAssignerState.Opening.TYPE), + @JsonSubTypes.Type( + value = ImmutableWaitingUntilCloseable.class, + name = BucketAssignerState.WaitingUntilCloseable.TYPE), + @JsonSubTypes.Type(value = ImmutableClosingFromOpen.class, name = BucketAssignerState.ClosingFromOpen.TYPE), + @JsonSubTypes.Type(value = ImmutableImmediatelyClosing.class, name = BucketAssignerState.ImmediatelyClosing.TYPE), +}) +public interface BucketAssignerState { + T accept(BucketAssignerState.Visitor visitor); + + static BucketAssignerState.Start start(long startTimestampInclusive) { + return ImmutableStart.of(startTimestampInclusive); + } + + static BucketAssignerState.Opening opening(long startTimestampInclusive) { + return ImmutableOpening.of(startTimestampInclusive); + } + + static BucketAssignerState.WaitingUntilCloseable waitingUntilCloseable(long startTimestampInclusive) { + return ImmutableWaitingUntilCloseable.of(startTimestampInclusive); + } + + static BucketAssignerState.ClosingFromOpen closingFromOpen( + long startTimestampInclusive, long endTimestampExclusive) { + return ImmutableClosingFromOpen.of(startTimestampInclusive, endTimestampExclusive); + } + + static BucketAssignerState.ImmediatelyClosing immediatelyClosing( + long startTimestampInclusive, long endTimestampExclusive) { + return ImmutableImmediatelyClosing.of(startTimestampInclusive, endTimestampExclusive); + } + + @Value.Immutable + @JsonSerialize(as = ImmutableStart.class) + @JsonDeserialize(as = ImmutableStart.class) + interface Start extends BucketAssignerState { + String TYPE = "start"; + + @Value.Parameter + long startTimestampInclusive(); + + @Override + default T accept(BucketAssignerState.Visitor visitor) { + return visitor.visit(this); + } + } + + @Value.Immutable + @JsonSerialize(as = ImmutableOpening.class) + @JsonDeserialize(as = ImmutableOpening.class) + interface Opening extends BucketAssignerState { + String TYPE = "opening"; + + @Value.Parameter + long startTimestampInclusive(); + + @Override + default T accept(BucketAssignerState.Visitor visitor) { + return visitor.visit(this); + } + } + + @Value.Immutable + @JsonSerialize(as = ImmutableWaitingUntilCloseable.class) + @JsonDeserialize(as = ImmutableWaitingUntilCloseable.class) + interface WaitingUntilCloseable extends BucketAssignerState { + String TYPE = "waitingUntilCloseable"; + + @Value.Parameter + long startTimestampInclusive(); + + @Override + default T accept(BucketAssignerState.Visitor visitor) { + return visitor.visit(this); + } + } + + @Value.Immutable + @JsonSerialize(as = ImmutableClosingFromOpen.class) + @JsonDeserialize(as = ImmutableClosingFromOpen.class) + interface ClosingFromOpen extends BucketAssignerState { + String TYPE = "closingFromOpen"; + + @Value.Parameter + long startTimestampInclusive(); + + @Value.Parameter + long endTimestampExclusive(); + + @Override + default T accept(BucketAssignerState.Visitor visitor) { + return visitor.visit(this); + } + } + + @Value.Immutable + @JsonSerialize(as = ImmutableImmediatelyClosing.class) + @JsonDeserialize(as = ImmutableImmediatelyClosing.class) + interface ImmediatelyClosing extends BucketAssignerState { + String TYPE = "immediatelyClosing"; + + @Value.Parameter + long startTimestampInclusive(); + + @Value.Parameter + long endTimestampExclusive(); + + @Override + default T accept(BucketAssignerState.Visitor visitor) { + return visitor.visit(this); + } + } + + interface Visitor { + T visit(BucketAssignerState.Start start); + + T visit(BucketAssignerState.Opening opening); + + T visit(BucketAssignerState.WaitingUntilCloseable waitingUntilCloseable); + + T visit(BucketAssignerState.ClosingFromOpen closingFromOpen); + + T visit(BucketAssignerState.ImmediatelyClosing immediatelyClosing); + } +} diff --git a/atlasdb-impl-shared/src/test/java/com/palantir/atlasdb/sweep/asts/bucketingthings/BucketAssignerStateTest.java b/atlasdb-impl-shared/src/test/java/com/palantir/atlasdb/sweep/asts/bucketingthings/BucketAssignerStateTest.java new file mode 100644 index 0000000000..2280c901f8 --- /dev/null +++ b/atlasdb-impl-shared/src/test/java/com/palantir/atlasdb/sweep/asts/bucketingthings/BucketAssignerStateTest.java @@ -0,0 +1,73 @@ +/* + * (c) Copyright 2024 Palantir Technologies Inc. All rights reserved. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +package com.palantir.atlasdb.sweep.asts.bucketingthings; + +import static org.assertj.core.api.Assertions.assertThat; + +import com.fasterxml.jackson.databind.ObjectMapper; +import com.google.common.io.BaseEncoding; +import com.palantir.conjure.java.serialization.ObjectMappers; +import java.io.IOException; +import java.util.stream.Stream; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.Arguments; +import org.junit.jupiter.params.provider.MethodSource; + +public class BucketAssignerStateTest { + private static final ObjectMapper OBJECT_MAPPER = ObjectMappers.newSmileServerObjectMapper(); + + // Think _very_ carefully about changing this without a migration. + private static final byte[] START_SERIALIZED = + BaseEncoding.base64().decode("OikKBfqDdHlwZURzdGFydJZzdGFydFRpbWVzdGFtcEluY2x1c2l2ZcL7"); + private static final byte[] OPENING_SERIALIZED = + BaseEncoding.base64().decode("OikKBfqDdHlwZUZvcGVuaW5nlnN0YXJ0VGltZXN0YW1wSW5jbHVzaXZlxPs="); + private static final byte[] WAITING_UNTIL_CLOSED_SERIALIZED = BaseEncoding.base64() + .decode("OikKBfqDdHlwZVR3YWl0aW5nVW50aWxDbG9zZWFibGWWc3RhcnRUaW1lc3RhbXBJbmNsdXNpdmXG+w=="); + private static final byte[] CLOSING_FROM_OPEN_SERIALIZED = BaseEncoding.base64() + .decode( + "OikKBfqDdHlwZU5jbG9zaW5nRnJvbU9wZW6Wc3RhcnRUaW1lc3RhbXBJbmNsdXNpdmXIlGVuZFRpbWVzdGFtcEV4Y2x1c2l2Zcr7"); + private static final byte[] IMMEDIATELY_CLOSING_SERIALIZED = BaseEncoding.base64() + .decode( + "OikKBfqDdHlwZVFpbW1lZGlhdGVseUNsb3NpbmeWc3RhcnRUaW1lc3RhbXBJbmNsdXNpdmXMlGVuZFRpbWVzdGFtcEV4Y2x1c2l2Zc77"); + + @ParameterizedTest(name = "{0}") + @MethodSource("bucketAssignerStates") + public void canRoundTripSerdeState(BucketAssignerState state, byte[] _unused) throws IOException { + byte[] json = OBJECT_MAPPER.writeValueAsBytes(state); + BucketAssignerState deserialized = OBJECT_MAPPER.readValue(json, BucketAssignerState.class); + assertThat(deserialized).isEqualTo(state); + // Adding instance of check to be _really_ sure, since a mistake here would be very painful. + assertThat(deserialized).isInstanceOf(state.getClass()); + } + + @ParameterizedTest(name = "{0}") + @MethodSource("bucketAssignerStates") + public void canDeserializeExistingVersion(BucketAssignerState bucketProgress, byte[] serializedForm) + throws IOException { + assertThat(OBJECT_MAPPER.readValue(serializedForm, BucketAssignerState.class)) + .isEqualTo(bucketProgress); + } + + private static Stream bucketAssignerStates() { + return Stream.of( + Arguments.of(BucketAssignerState.start(1), START_SERIALIZED), + Arguments.of(BucketAssignerState.opening(2), OPENING_SERIALIZED), + Arguments.of(BucketAssignerState.waitingUntilCloseable(3), WAITING_UNTIL_CLOSED_SERIALIZED), + Arguments.of(BucketAssignerState.closingFromOpen(4, 5), CLOSING_FROM_OPEN_SERIALIZED), + Arguments.of(BucketAssignerState.immediatelyClosing(6, 7), IMMEDIATELY_CLOSING_SERIALIZED)); + } +}