@Override
    public void put(final Bytes rawBaseKey,
                    final byte[] value) {
        final long timestamp = baseKeySchema.segmentTimestamp(rawBaseKey);
        observedStreamTime = Math.max(observedStreamTime, timestamp);
        final long segmentId = segments.segmentId(timestamp);
        final S segment = segments.getOrCreateSegmentIfLive(segmentId, context, observedStreamTime);

        if (segment == null) {
            expiredRecordSensor.record(1.0d, context.currentSystemTimeMs());
            LOG.warn("Skipping record for expired segment.");
        } else {
            synchronized (position) {
                StoreQueryUtils.updatePosition(position, stateStoreContext);

                // Put to index first so that if put to base failed, when we iterate index, we will
                // find no base value. If put to base first but putting to index fails, when we iterate
                // index, we can't find the key but if we iterate over base store, we can find the key
                // which lead to inconsistency.
                if (hasIndex()) {
                    final KeyValue<Bytes, byte[]> indexKeyValue = getIndexKeyValue(rawBaseKey, value);
                    segment.put(indexKeyValue.key, indexKeyValue.value);
                }
                segment.put(rawBaseKey, value);
            }
        }
    }
@Test
    public void shouldFetchSessionForSingleKey() {
        // Only for TimeFirstSessionKeySchema schema
        if (!(getBaseSchema() instanceof TimeFirstSessionKeySchema)) {
            return;
        }

        final String keyA = "a";
        final String keyB = "b";
        final String keyC = "c";

        final StateSerdes<String, Long> stateSerdes = StateSerdes.withBuiltinTypes("dummy", String.class, Long.class);
        final Bytes key1 = Bytes.wrap(stateSerdes.keySerializer().serialize("dummy", keyA));
        final Bytes key2 = Bytes.wrap(stateSerdes.keySerializer().serialize("dummy", keyB));
        final Bytes key3 = Bytes.wrap(stateSerdes.keySerializer().serialize("dummy", keyC));

        final byte[] expectedValue1 = serializeValue(10);
        final byte[] expectedValue2 = serializeValue(50);
        final byte[] expectedValue3 = serializeValue(100);
        final byte[] expectedValue4 = serializeValue(200);

        bytesStore.put(serializeKey(new Windowed<>(keyA, windows[0])), expectedValue1);
        bytesStore.put(serializeKey(new Windowed<>(keyA, windows[1])), expectedValue2);
        bytesStore.put(serializeKey(new Windowed<>(keyB, windows[2])), expectedValue3);
        bytesStore.put(serializeKey(new Windowed<>(keyC, windows[3])), expectedValue4);

        // Record expired as timestampFromRawKey = 1000 while observedStreamTime = 60,000 and retention = 1000.
        final byte[] value1 = ((RocksDBTimeOrderedSessionSegmentedBytesStore) bytesStore).fetchSession(
            key1, windows[0].start(), windows[0].end());
        assertNull(value1);

        // Record expired as timestampFromRawKey = 1000 while observedStreamTime = 60,000 and retention = 1000.
        final byte[] value2 = ((RocksDBTimeOrderedSessionSegmentedBytesStore) bytesStore).fetchSession(
            key1, windows[1].start(), windows[1].end());
        assertNull(value2);

        // expired record
        // timestampFromRawKey = 1500 while observedStreamTime = 60,000 and retention = 1000.
        final byte[] value3 = ((RocksDBTimeOrderedSessionSegmentedBytesStore) bytesStore).fetchSession(
            key2, windows[2].start(), windows[2].end());
        assertNull(value3);

        // only non-expired record
        // timestampFromRawKey = 60,000 while observedStreamTime = 60,000 and retention = 1000.
        final byte[] value4 = ((RocksDBTimeOrderedSessionSegmentedBytesStore) bytesStore).fetchSession(
            key3, windows[3].start(), windows[3].end());
        assertEquals(Bytes.wrap(value4), Bytes.wrap(expectedValue4));

        final byte[] noValue = ((RocksDBTimeOrderedSessionSegmentedBytesStore) bytesStore).fetchSession(
            key3, 2000, 3000);
        assertNull(noValue);
    }