<commit_before>/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 org.apache.jackrabbit.oak.plugins.segment;

import static com.google.common.base.Objects.equal;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkElementIndex;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkPositionIndex;
import static com.google.common.base.Preconditions.checkPositionIndexes;
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.collect.Lists.newArrayList;
import static java.util.Collections.emptyMap;
import static org.apache.jackrabbit.oak.plugins.segment.MapRecord.BUCKETS_PER_LEVEL;
import static org.apache.jackrabbit.oak.plugins.segment.Segment.MAX_SEGMENT_SIZE;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.UUID;

import javax.jcr.PropertyType;

import org.apache.jackrabbit.oak.api.Blob;
import org.apache.jackrabbit.oak.api.PropertyState;
import org.apache.jackrabbit.oak.api.Type;
import org.apache.jackrabbit.oak.plugins.memory.ModifiedNodeState;
import org.apache.jackrabbit.oak.spi.state.ChildNodeEntry;
import org.apache.jackrabbit.oak.spi.state.DefaultNodeStateDiff;
import org.apache.jackrabbit.oak.spi.state.NodeState;

import com.google.common.base.Charsets;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.io.ByteStreams;

public class SegmentWriter {

    static final int BLOCK_SIZE = 1 << 12; // 4kB

    private final SegmentStore store;

    private final Map<String, RecordId> strings = new LinkedHashMap<String, RecordId>(1500, 0.75f, true) {
        @Override
        protected boolean removeEldestEntry(Entry<String, RecordId> eldest) {
            return size() > 1000;
        }
    };

    private final Map<Template, RecordId> templates = new LinkedHashMap<Template, RecordId>(1500, 0.75f, true) {
        @Override
        protected boolean removeEldestEntry(Entry<Template, RecordId> eldest) {
            return size() > 1000;
        }
    };

    private UUID uuid = UUID.randomUUID();

    /**
     * Insertion-ordered map from the UUIDs of referenced segments to the
     * respective single-byte UUID index values used when serializing
     * record identifiers.
     */
    private final Map<UUID, Byte> uuids = Maps.newLinkedHashMap();

    /**
     * The segment write buffer, filled from the end to the beginning
     * (see OAK-629). The buffer grows automatically up to
     * {@link Segment#MAX_SEGMENT_SIZE}.
     */
    private byte[] buffer = new byte[MAX_SEGMENT_SIZE];

    /**
     * The number of bytes already written (or allocated). Counted from
     * the <em>end</em> of the buffer.
     */
    private int length = 0;

    /**
     * Current write position within the buffer. Grows up when raw data
     * is written, but shifted downwards by the prepare methods.
     */
    private int position;

    public SegmentWriter(SegmentStore store) {
        this.store = store;
    }

    public synchronized Segment getCurrentSegment(UUID id) {
        if (equal(id, uuid)) {
            return new Segment(
                    store, uuid,
                    ByteBuffer.wrap(buffer, buffer.length - length, length),
                    newArrayList(uuids.keySet()));
        } else {
            return null;
        }
    }

    public synchronized void flush() {
        if (length > 0) {
            store.writeSegment(
                    uuid, buffer, buffer.length - length, length,
                    newArrayList(uuids.keySet()));

            uuid = UUID.randomUUID();
            uuids.clear();
            buffer = new byte[MAX_SEGMENT_SIZE];
            length = 0;
            position = buffer.length;
        }
    }

    private RecordId prepare(int size) {
        return prepare(size, Collections.<RecordId>emptyList());
    }

    private RecordId prepare(int size, Collection<RecordId> ids) {
        checkArgument(size >= 0);

        Set<UUID> segmentIds = new HashSet<UUID>();
        for (RecordId id : checkNotNull(ids)) {
            UUID segmentId = id.getSegmentId();
            if (!uuids.containsKey(segmentId)) {
                segmentIds.add(segmentId);
            }
        }

        int fullSize = size + ids.size() * Segment.RECORD_ID_BYTES;
        checkArgument(fullSize > 0);

        int alignment = Segment.RECORD_ALIGN_BYTES - 1;
        int alignedSize = (fullSize + alignment) & ~alignment;
        if (length + alignedSize > buffer.length
                || uuids.size() + segmentIds.size() > 0x100) {
            flush();
        }

        length += alignedSize;
        position = buffer.length - length;
        checkState(position >= 0);
        return new RecordId(uuid, position);
    }

    private void writeRecordId(RecordId id) {
        checkNotNull(id);

        UUID segmentId = id.getSegmentId();
        Byte segmentIndex = uuids.get(segmentId);
        if (segmentIndex == null) {
            checkState(uuids.size() < Segment.SEGMENT_REFERENCE_LIMIT);
            segmentIndex = Byte.valueOf((byte) uuids.size());
            uuids.put(segmentId, segmentIndex);
        }

        int offset = id.getOffset();
        checkState(0 <= offset && offset < MAX_SEGMENT_SIZE);
        checkState((offset & (Segment.RECORD_ALIGN_BYTES - 1)) == 0);

        buffer[position++] = segmentIndex.byteValue();<commit_msg>OAK-1032: SegmentMK: One SegmentWriter per SegmentNodeStore<commit_after>/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 org.apache.jackrabbit.oak.plugins.segment;

import static com.google.common.base.Objects.equal;
import static com.google.common.base.Preconditions.checkArgument;
import static com.google.common.base.Preconditions.checkElementIndex;
import static com.google.common.base.Preconditions.checkNotNull;
import static com.google.common.base.Preconditions.checkPositionIndex;
import static com.google.common.base.Preconditions.checkPositionIndexes;
import static com.google.common.base.Preconditions.checkState;
import static com.google.common.collect.Lists.newArrayList;
import static java.util.Collections.emptyMap;
import static org.apache.jackrabbit.oak.plugins.segment.MapRecord.BUCKETS_PER_LEVEL;
import static org.apache.jackrabbit.oak.plugins.segment.Segment.MAX_SEGMENT_SIZE;

import java.io.ByteArrayInputStream;
import java.io.IOException;
import java.io.InputStream;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collection;
import java.util.Collections;
import java.util.HashSet;
import java.util.LinkedHashMap;
import java.util.List;
import java.util.Map;
import java.util.Map.Entry;
import java.util.Set;
import java.util.UUID;

import javax.jcr.PropertyType;

import org.apache.jackrabbit.oak.api.Blob;
import org.apache.jackrabbit.oak.api.PropertyState;
import org.apache.jackrabbit.oak.api.Type;
import org.apache.jackrabbit.oak.plugins.memory.ModifiedNodeState;
import org.apache.jackrabbit.oak.spi.state.ChildNodeEntry;
import org.apache.jackrabbit.oak.spi.state.DefaultNodeStateDiff;
import org.apache.jackrabbit.oak.spi.state.NodeState;

import com.google.common.base.Charsets;
import com.google.common.collect.Iterables;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.io.ByteStreams;

public class SegmentWriter {

    static final int BLOCK_SIZE = 1 << 12; // 4kB

    private final SegmentStore store;

    private final Map<String, RecordId> strings = new LinkedHashMap<String, RecordId>(1500, 0.75f, true) {
        @Override
        protected boolean removeEldestEntry(Entry<String, RecordId> eldest) {
            return size() > 1000;
        }
    };

    private final Map<Template, RecordId> templates = new LinkedHashMap<Template, RecordId>(1500, 0.75f, true) {
        @Override
        protected boolean removeEldestEntry(Entry<Template, RecordId> eldest) {
            return size() > 1000;
        }
    };

    private UUID uuid = UUID.randomUUID();

    /**
     * Insertion-ordered map from the UUIDs of referenced segments to the
     * respective single-byte UUID index values used when serializing
     * record identifiers.
     */
    private final Map<UUID, Byte> uuids = Maps.newLinkedHashMap();

    /**
     * The segment write buffer, filled from the end to the beginning
     * (see OAK-629).
     */
    private final byte[] buffer = new byte[MAX_SEGMENT_SIZE];

    /**
     * The number of bytes already written (or allocated). Counted from
     * the <em>end</em> of the buffer.
     */
    private int length = 0;

    /**
     * Current write position within the buffer. Grows up when raw data
     * is written, but shifted downwards by the prepare methods.
     */
    private int position;

    private Segment currentSegment = null;

    public SegmentWriter(SegmentStore store) {
        this.store = store;
    }

    public synchronized Segment getCurrentSegment(UUID id) {
        if (equal(id, uuid)) {
            if (currentSegment == null) {
                ByteBuffer b = ByteBuffer.allocate(length);
                System.arraycopy(
                        buffer, buffer.length - length, b.array(), 0, length);
                currentSegment = new Segment(
                        store, uuid, b, newArrayList(uuids.keySet()));
            }
            return currentSegment;
        } else {
            return null;
        }
    }

    public synchronized void flush() {
        if (length > 0) {
            store.writeSegment(
                    uuid, buffer, buffer.length - length, length,
                    newArrayList(uuids.keySet()));

            uuid = UUID.randomUUID();
            uuids.clear();
            length = 0;
            position = buffer.length;
            currentSegment = null;
        }
    }

    private RecordId prepare(int size) {
        return prepare(size, Collections.<RecordId>emptyList());
    }

    private RecordId prepare(int size, Collection<RecordId> ids) {
        checkArgument(size >= 0);

        Set<UUID> segmentIds = new HashSet<UUID>();
        for (RecordId id : checkNotNull(ids)) {
            UUID segmentId = id.getSegmentId();
            if (!uuids.containsKey(segmentId)) {
                segmentIds.add(segmentId);
            }
        }

        int fullSize = size + ids.size() * Segment.RECORD_ID_BYTES;
        checkArgument(fullSize > 0);

        int alignment = Segment.RECORD_ALIGN_BYTES - 1;
        int alignedSize = (fullSize + alignment) & ~alignment;
        if (length + alignedSize > buffer.length
                || uuids.size() + segmentIds.size() > 0x100) {
            flush();
        }

        length += alignedSize;
        position = buffer.length - length;
        checkState(position >= 0);
        currentSegment = null;
        return new RecordId(uuid, position);
    }

    private void writeRecordId(RecordId id) {
        checkNotNull(id);

        UUID segmentId = id.getSegmentId();
        Byte segmentIndex = uuids.get(segmentId);
        if (segmentIndex == null) {
            checkState(uuids.size() < Segment.SEGMENT_REFERENCE_LIMIT);
            segmentIndex = Byte.valueOf((byte) uuids.size());
            uuids.put(segmentId, segmentIndex);
        }

        int offset = id.getOffset();
        checkState(0 <= offset && offset < MAX_SEGMENT_SIZE);
        checkState((offset & (Segment.RECORD_ALIGN_BYTES - 1)) == 0);

        buffer[position++] = segmentIndex.byteValue();
{'fraction_non_alphanumeric': 0.09898345329079172, 'fraction_numerical': 0.0073458484825999855, 'mean_word_length': 2.8552631578947367, 'pattern_counts': {'":': 0, '<': 39, '<?xml version=': 0, '>': 45, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5329198', 'n_tokens_mistral': 3867, 'n_tokens_neox': 3718, 'n_words': 1018}