Java tutorial
/* * 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 * * 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.streams.mongo; import org.apache.streams.config.ComponentConfigurator; import org.apache.streams.config.StreamsConfigurator; import org.apache.streams.core.DatumStatusCounter; import org.apache.streams.core.StreamsDatum; import org.apache.streams.core.StreamsPersistReader; import org.apache.streams.core.StreamsResultSet; import org.apache.streams.jackson.StreamsJacksonMapper; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.node.ObjectNode; import com.google.common.collect.Queues; import com.mongodb.DB; import com.mongodb.DBCollection; import com.mongodb.DBCursor; import com.mongodb.DBObject; import com.mongodb.MongoClient; import com.mongodb.MongoCredential; import com.mongodb.ServerAddress; import org.apache.commons.lang3.StringUtils; import org.joda.time.DateTime; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.IOException; import java.math.BigInteger; import java.util.Queue; import java.util.concurrent.ExecutorService; import java.util.concurrent.Executors; import java.util.concurrent.Future; import java.util.concurrent.LinkedBlockingQueue; import java.util.concurrent.atomic.AtomicLong; import java.util.concurrent.locks.ReadWriteLock; import java.util.concurrent.locks.ReentrantReadWriteLock; import java.util.stream.Collectors; import java.util.stream.Stream; /** * MongoPersistReader reads documents from mongo. */ public class MongoPersistReader implements StreamsPersistReader { public static final String STREAMS_ID = "MongoPersistReader"; private static final Logger LOGGER = LoggerFactory.getLogger(MongoPersistReader.class); protected volatile Queue<StreamsDatum> persistQueue; private ObjectMapper mapper = StreamsJacksonMapper.getInstance(); private volatile AtomicLong lastWrite = new AtomicLong(System.currentTimeMillis()); private ExecutorService executor; private MongoConfiguration config; protected MongoClient client; protected DB db; protected DBCollection collection; protected DBCursor cursor; protected final ReadWriteLock lock = new ReentrantReadWriteLock(); /** * KafkaPersistReader constructor - resolves KafkaConfiguration from JVM 'mongo'. */ public MongoPersistReader() { this.config = new ComponentConfigurator<>(MongoConfiguration.class) .detectConfiguration(StreamsConfigurator.getConfig().getConfig("mongo")); } /** * KafkaPersistReader constructor - uses supplied MongoConfiguration. * @param config config */ public MongoPersistReader(MongoConfiguration config) { this.config = config; } /** * KafkaPersistReader constructor - uses supplied persistQueue. * @param persistQueue persistQueue */ public MongoPersistReader(Queue<StreamsDatum> persistQueue) { this.config = new ComponentConfigurator<>(MongoConfiguration.class) .detectConfiguration(StreamsConfigurator.getConfig().getConfig("mongo")); this.persistQueue = persistQueue; } public void setPersistQueue(Queue<StreamsDatum> persistQueue) { this.persistQueue = persistQueue; } public Queue<StreamsDatum> getPersistQueue() { return persistQueue; } public void stop() { } @Override public String getId() { return STREAMS_ID; } @Override public void prepare(Object configurationObject) { connectToMongo(); if (client == null || collection == null) { throw new RuntimeException("Unable to connect!"); } cursor = collection.find(); if (cursor == null || !cursor.hasNext()) { throw new RuntimeException("Collection not present or empty!"); } persistQueue = constructQueue(); executor = Executors.newSingleThreadExecutor(); } @Override public void cleanUp() { stop(); } protected StreamsDatum prepareDatum(DBObject dbObject) { ObjectNode objectNode; String id; try { objectNode = mapper.readValue(dbObject.toString(), ObjectNode.class); id = objectNode.get("_id").get("$oid").asText(); objectNode.remove("_id"); } catch (IOException ex) { LOGGER.warn("document isn't valid JSON."); return null; } return new StreamsDatum(objectNode, id); } private synchronized void connectToMongo() { ServerAddress serverAddress = new ServerAddress(config.getHost(), config.getPort().intValue()); if (StringUtils.isNotEmpty(config.getUser()) && StringUtils.isNotEmpty(config.getPassword())) { MongoCredential credential = MongoCredential.createCredential(config.getUser(), config.getDb(), config.getPassword().toCharArray()); client = new MongoClient(serverAddress, Stream.of(credential).collect(Collectors.toList())); } else { client = new MongoClient(serverAddress); } db = client.getDB(config.getDb()); if (!db.collectionExists(config.getCollection())) { db.createCollection(config.getCollection(), null); } collection = db.getCollection(config.getCollection()); } @Override public StreamsResultSet readAll() { try (DBCursor cursor = collection.find()) { while (cursor.hasNext()) { DBObject dbObject = cursor.next(); StreamsDatum datum = prepareDatum(dbObject); write(datum); } } return readCurrent(); } @Override public void startStream() { LOGGER.debug("startStream"); MongoPersistReaderTask readerTask = new MongoPersistReaderTask(this); Thread readerTaskThread = new Thread(readerTask); Future future = executor.submit(readerTaskThread); while (!future.isDone() && !future.isCancelled()) { try { Thread.sleep(1000); } catch (InterruptedException interrupt) { LOGGER.trace("Interrupt", interrupt); } } executor.shutdown(); } @Override public StreamsResultSet readCurrent() { StreamsResultSet current; try { lock.writeLock().lock(); current = new StreamsResultSet(persistQueue); current.setCounter(new DatumStatusCounter()); persistQueue = constructQueue(); } finally { lock.writeLock().unlock(); } return current; } //The locking may appear to be counter intuitive but we really don't care if multiple threads offer to the queue //as it is a synchronized queue. What we do care about is that we don't want to be offering to the current reference //if the queue is being replaced with a new instance protected void write(StreamsDatum entry) { boolean success; do { try { lock.readLock().lock(); success = persistQueue.offer(entry); Thread.yield(); } finally { lock.readLock().unlock(); } } while (!success); } @Override public StreamsResultSet readNew(BigInteger sequence) { return null; } @Override public StreamsResultSet readRange(DateTime start, DateTime end) { return null; } @Override public boolean isRunning() { return !executor.isTerminated() || !executor.isShutdown(); } private Queue<StreamsDatum> constructQueue() { return Queues.synchronizedQueue(new LinkedBlockingQueue<StreamsDatum>(10000)); } public class MongoPersistReaderTask implements Runnable { private MongoPersistReader reader; public MongoPersistReaderTask(MongoPersistReader reader) { this.reader = reader; } @Override public void run() { try { while (reader.cursor.hasNext()) { DBObject dbObject = reader.cursor.next(); StreamsDatum datum = reader.prepareDatum(dbObject); reader.write(datum); } } finally { reader.cursor.close(); } } } }