com.crosstreelabs.cognitio.service.mongo.MongoHostService.java Source code

Introduction

Here is the source code for com.crosstreelabs.cognitio.service.mongo.MongoHostService.java

Source

/*
 * Copyright 2015 Crosstree Labs.
 *
 * 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.crosstreelabs.cognitio.service.mongo;

import com.crosstreelabs.cognitio.api.resource.Host;
import com.crosstreelabs.cognitio.api.service.HostService;
import com.mongodb.BasicDBObject;
import com.mongodb.DB;
import com.mongodb.DBCollection;
import com.mongodb.DBCursor;
import com.mongodb.DBObject;
import com.mongodb.MongoClient;
import com.mongodb.WriteConcern;
import java.util.ArrayList;
import java.util.Collection;
import java.util.Collections;
import java.util.Date;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Objects;
import java.util.Queue;
import org.apache.commons.collections4.queue.CircularFifoQueue;
import org.apache.commons.lang3.StringUtils;
import org.bson.types.ObjectId;
import org.joda.time.DateTime;

public class MongoHostService implements HostService {
    private final MongoClient client;
    private final DB db;
    private final DBCollection hosts;

    private final Queue<CacheEntry> hostCache = new CircularFifoQueue<>(50);

    public MongoHostService(final MongoClient client) {
        this.client = client;
        db = client.getDB("cognitio");
        if (!db.collectionExists("hosts")) {
            hosts = db.createCollection("hosts", null);
            hosts.createIndex(new BasicDBObject("host", 1), new BasicDBObject("unique", true));
        } else {
            hosts = db.getCollection("hosts");
        }
    }

    @Override
    public Host findByDomain(final String domain) {
        for (CacheEntry entry : hostCache) {
            if (entry.getHost().host.equals(domain)) {
                return entry.getHost();
            }
        }

        DBObject obj = hosts.findOne(new BasicDBObject("host", domain));
        if (obj == null) {
            return null;
        }
        Host host = fromMongoObject(obj);
        hostCache.add(new CacheEntry(host));
        return host;
    }

    @Override
    public Host findById(final String id) {
        for (CacheEntry entry : hostCache) {
            if (entry.getHost().id.equals(id)) {
                return entry.getHost();
            }
        }

        DBObject obj = hosts.findOne(new BasicDBObject("_id", new ObjectId(id)));
        if (obj == null) {
            return null;
        }
        Host host = fromMongoObject(obj);
        hostCache.add(new CacheEntry(host));
        return host;
    }

    @Override
    public Collection<Host> recentlyRequested(final int withinSeconds) {
        DBCursor hostsCursor = hosts.find(new BasicDBObject("last_request",
                new BasicDBObject("$gte", new DateTime().minusSeconds(withinSeconds).toDate())));
        List<Host> result = new ArrayList<>();
        for (DBObject obj : hostsCursor) {
            result.add(fromMongoObject(obj));
        }
        return result;
    }

    @Override
    public void save(Host host) {
        if (host == null) {
            throw new IllegalArgumentException("Host cannot be null");
        }
        if (StringUtils.isBlank(host.host)) {
            throw new IllegalArgumentException("Cannot save host with empty domain");
        }

        // If the host is cached, let's check to see if it changed
        for (CacheEntry entry : hostCache) {
            if (entry.getHost().equals(host) && !entry.isDirty()) {
                return;
            }
        }

        // Save the entry
        if (host.id == null) {
            DBObject obj = toMongoObject(host);
            hosts.insert(obj, WriteConcern.JOURNALED);
            host.id = obj.get("_id").toString();
        } else {
            hosts.update(new BasicDBObject("_id", new ObjectId(host.id)), toMongoObject(host), false, false,
                    WriteConcern.JOURNALED);
        }

        // Add to cache if needed
        for (CacheEntry entry : hostCache) {
            if (entry.getHost().equals(host)) {
                return;
            }
        }
        hostCache.add(new CacheEntry(host));
    }

    //~ Conversion methods ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
    protected DBObject toMongoObject(final Host host) {
        BasicDBObject result = new BasicDBObject();
        result.append("host", host.host);
        if (host.lastRequest != null) {
            result.append("last_request", host.lastRequest.toDate());
        }
        result.append("robots_txt", host.robotsTxt);
        return result;
    }

    protected Host fromMongoObject(final DBObject obj) {
        Host result = new Host();
        result.id = obj.get("_id").toString();
        result.host = obj.get("host").toString();
        if (obj.get("last_request") != null) {
            result.lastRequest = new DateTime((Date) obj.get("last_request"));
        }
        if (obj.get("robots_txt") != null) {
            result.robotsTxt = obj.get("robots_txt").toString();
        }
        return result;
    }

    private static final class CacheEntry {
        private final Host host;
        private Map<String, Object> originalData;

        public CacheEntry(final Host host) {
            this.host = host;
            update();
        }

        public Host getHost() {
            return host;
        }

        public boolean isDirty() {
            return !(Objects.equals(host.id, originalData.get("id"))
                    && Objects.equals(host.host, originalData.get("host"))
                    && Objects.equals(host.lastRequest, originalData.get("lastRequest"))
                    && Objects.equals(host.robotsTxt, originalData.get("robotsTxt")));
        }

        public void update() {
            Map<String, Object> data = new HashMap<>();
            data.put("id", host.id);
            data.put("host", host.host);
            data.put("lastRequest", host.lastRequest);
            data.put("robotsTxt", host.robotsTxt);
            this.originalData = Collections.unmodifiableMap(data);
        }
    }

}