<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 java.util.List;
import java.util.UUID;
import java.util.concurrent.ExecutionException;

import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;
import com.mongodb.BasicDBObject;
import com.mongodb.DB;
import com.mongodb.DBCollection;
import com.mongodb.DBObject;
import com.mongodb.Mongo;

public class MongoStore implements SegmentStore {

    private static final int MAX_SEGMENT_SIZE = 1 << 23; // 8MB

    private static final long DEFAULT_CACHE_SIZE = 1 << 28; // 256MB

    private final DBCollection segments;

    private final DBCollection journals;

    private final LoadingCache<UUID, Segment> cache;

    public MongoStore(DB db, long cacheSize) {
        this.segments = db.getCollection("segments");
        this.journals = db.getCollection("journals");

        this.cache = CacheBuilder.newBuilder()
                .maximumWeight(cacheSize)
                .weigher(Segment.weigher())
                .build(new CacheLoader<UUID, Segment>() {
                    @Override
                    public Segment load(UUID key) throws Exception {
                        return findSegment(key);
                    }
                });
    }

    public MongoStore(DB db) {
        this(db, DEFAULT_CACHE_SIZE);
    }


    public MongoStore(Mongo mongo) {
        this(mongo.getDB("Oak"));
    }

    @Override
    public RecordId getJournalHead() {
        DBObject journal = journals.findOne(new BasicDBObject("_id", "root"));
        return RecordId.fromString(journal.get("head").toString());
    }

    @Override
    public boolean setJournalHead(RecordId head, RecordId base) {
        DBObject baseObject = new BasicDBObject(
                ImmutableMap.of("_id", "root", "head", base.toString()));
        DBObject headObject = new BasicDBObject(
                ImmutableMap.of("_id", "root", "head", head.toString()));
        return journals.update(baseObject, headObject).getN() == 1;
    }

    @Override
    public int getMaxSegmentSize() {
        return MAX_SEGMENT_SIZE;
    }

    @Override
    public Segment readSegment(UUID segmentId) {
        try {
            return cache.get(segmentId);
        } catch (ExecutionException e) {
            throw new IllegalStateException(
                    "Failed to read segment " + segmentId, e);
        }
    }

    @Override
    public void createSegment(Segment segment) {<commit_msg>OAK-593: Segment-based MK<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 java.util.List;
import java.util.UUID;
import java.util.concurrent.ExecutionException;

import org.apache.jackrabbit.oak.plugins.memory.MemoryNodeState;

import com.google.common.cache.CacheBuilder;
import com.google.common.cache.CacheLoader;
import com.google.common.cache.LoadingCache;
import com.google.common.collect.ImmutableMap;
import com.google.common.collect.Lists;
import com.mongodb.BasicDBObject;
import com.mongodb.DB;
import com.mongodb.DBCollection;
import com.mongodb.DBObject;
import com.mongodb.Mongo;

public class MongoStore implements SegmentStore {

    private static final int MAX_SEGMENT_SIZE = 1 << 23; // 8MB

    private static final long DEFAULT_CACHE_SIZE = 1 << 28; // 256MB

    private final DBCollection segments;

    private final DBCollection journals;

    private final LoadingCache<UUID, Segment> cache;

    public MongoStore(DB db, long cacheSize) {
        this.segments = db.getCollection("segments");
        this.journals = db.getCollection("journals");

        this.cache = CacheBuilder.newBuilder()
                .maximumWeight(cacheSize)
                .weigher(Segment.weigher())
                .build(new CacheLoader<UUID, Segment>() {
                    @Override
                    public Segment load(UUID key) throws Exception {
                        return findSegment(key);
                    }
                });

        SegmentWriter writer = new SegmentWriter(this);
        RecordId id = writer.writeNode(MemoryNodeState.EMPTY_NODE);
        writer.flush();
        journals.update(
                new BasicDBObject("_id", "root"),
                new BasicDBObject(ImmutableMap.of(
                        "_id", "root", "head", id.toString())),
                true /* upsert */, false);
    }

    public MongoStore(DB db) {
        this(db, DEFAULT_CACHE_SIZE);
    }


    public MongoStore(Mongo mongo) {
        this(mongo.getDB("Oak"));
    }

    @Override
    public RecordId getJournalHead() {
        DBObject journal = journals.findOne(new BasicDBObject("_id", "root"));
        return RecordId.fromString(journal.get("head").toString());
    }

    @Override
    public boolean setJournalHead(RecordId head, RecordId base) {
        DBObject baseObject = new BasicDBObject(
                ImmutableMap.of("_id", "root", "head", base.toString()));
        DBObject headObject = new BasicDBObject(
                ImmutableMap.of("_id", "root", "head", head.toString()));
        return journals.findAndModify(baseObject, headObject) != null;
    }

    @Override
    public int getMaxSegmentSize() {
        return MAX_SEGMENT_SIZE;
    }

    @Override
    public Segment readSegment(UUID segmentId) {
        try {
            return cache.get(segmentId);
        } catch (ExecutionException e) {
            throw new IllegalStateException(
                    "Failed to read segment " + segmentId, e);
        }
    }

    @Override
    public void createSegment(Segment segment) {
{'fraction_non_alphanumeric': 0.09359132702072183, 'fraction_numerical': 0.004391381912995746, 'mean_word_length': 2.7703052250387996, 'pattern_counts': {'":': 0, '<': 15, '<?xml version=': 0, '>': 7, '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': '379354', 'n_tokens_mistral': 2016, 'n_tokens_neox': 1904, 'n_words': 581}