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.StreamsDatum; import org.apache.streams.core.StreamsPersistWriter; import org.apache.streams.jackson.StreamsJacksonMapper; import com.fasterxml.jackson.databind.ObjectMapper; import com.fasterxml.jackson.databind.node.ObjectNode; import com.mongodb.DB; import com.mongodb.DBCollection; import com.mongodb.DBObject; import com.mongodb.MongoClient; import com.mongodb.MongoCredential; import com.mongodb.ServerAddress; import com.mongodb.util.JSON; import org.apache.commons.lang3.StringUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; import java.io.Closeable; import java.io.Flushable; import java.io.IOException; import java.util.ArrayList; import java.util.List; import java.util.Queue; import java.util.Random; import java.util.concurrent.ConcurrentLinkedQueue; import java.util.concurrent.Executors; import java.util.concurrent.ScheduledExecutorService; import java.util.concurrent.TimeUnit; 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; public class MongoPersistWriter implements StreamsPersistWriter, Runnable, Flushable, Closeable { public static final String STREAMS_ID = "MongoPersistWriter"; private static final Logger LOGGER = LoggerFactory.getLogger(MongoPersistWriter.class); private static final long MAX_WRITE_LATENCY = 1000; protected volatile Queue<StreamsDatum> persistQueue; private ObjectMapper mapper = StreamsJacksonMapper.getInstance(); private volatile AtomicLong lastWrite = new AtomicLong(System.currentTimeMillis()); private ScheduledExecutorService backgroundFlushTask = Executors.newSingleThreadScheduledExecutor(); private MongoConfiguration config; protected MongoClient client; protected DB db; protected DBCollection collection; protected List<DBObject> insertBatch = new ArrayList<>(); protected final ReadWriteLock lock = new ReentrantReadWriteLock(); public MongoPersistWriter() { this(new ComponentConfigurator<>(MongoConfiguration.class) .detectConfiguration(StreamsConfigurator.getConfig().getConfig("mongo"))); } public MongoPersistWriter(MongoConfiguration config) { this.config = config; } public void setPersistQueue(Queue<StreamsDatum> persistQueue) { this.persistQueue = persistQueue; } public Queue<StreamsDatum> getPersistQueue() { return persistQueue; } @Override public String getId() { return STREAMS_ID; } @Override public void write(StreamsDatum streamsDatum) { DBObject dbObject = prepareObject(streamsDatum); if (dbObject != null) { addToBatch(dbObject); flushIfNecessary(); } } @Override public void flush() throws IOException { try { LOGGER.debug("Attempting to flush {} items to mongo", insertBatch.size()); lock.writeLock().lock(); collection.insert(insertBatch); lastWrite.set(System.currentTimeMillis()); insertBatch = new ArrayList<>(); } finally { lock.writeLock().unlock(); } } public synchronized void close() throws IOException { client.close(); backgroundFlushTask.shutdownNow(); } /** * start write thread. */ public void start() { connectToMongo(); backgroundFlushTask.scheduleAtFixedRate(this::flushIfNecessary, 0, MAX_WRITE_LATENCY * 2, TimeUnit.MILLISECONDS); } /** * stop. */ public void stop() { try { flush(); } catch (IOException ex) { LOGGER.error("Error flushing", ex); } try { close(); } catch (IOException ex) { LOGGER.error("Error closing", ex); } try { backgroundFlushTask.shutdown(); // Wait a while for existing tasks to terminate if (!backgroundFlushTask.awaitTermination(15, TimeUnit.SECONDS)) { backgroundFlushTask.shutdownNow(); // Cancel currently executing tasks // Wait a while for tasks to respond to being cancelled if (!backgroundFlushTask.awaitTermination(15, TimeUnit.SECONDS)) { LOGGER.error("Stream did not terminate"); } } } catch (InterruptedException ie) { // (Re-)Cancel if current thread also interrupted backgroundFlushTask.shutdownNow(); // Preserve interrupt status Thread.currentThread().interrupt(); } } @Override public void run() { while (true) { if (persistQueue.peek() != null) { try { StreamsDatum entry = persistQueue.remove(); write(entry); } catch (Exception ex) { ex.printStackTrace(); } } try { Thread.sleep(new Random().nextInt(1)); } catch (InterruptedException interrupt) { LOGGER.trace("Interrupt", interrupt); } } } @Override public void prepare(Object configurationObject) { this.persistQueue = new ConcurrentLinkedQueue<>(); start(); } @Override public void cleanUp() { stop(); } protected void flushIfNecessary() { long lastLatency = System.currentTimeMillis() - lastWrite.get(); //Flush iff the size > 0 AND the size is divisible by 100 or the time between now and the last flush is greater //than the maximum desired latency if (insertBatch.size() > 0 && (insertBatch.size() % 100 == 0 || lastLatency > MAX_WRITE_LATENCY)) { try { flush(); } catch (IOException ex) { LOGGER.error("Error writing to Mongo", ex); } } } protected void addToBatch(DBObject dbObject) { try { lock.readLock().lock(); insertBatch.add(dbObject); } finally { lock.readLock().unlock(); } } protected DBObject prepareObject(StreamsDatum streamsDatum) { DBObject dbObject = null; if (streamsDatum.getDocument() instanceof String) { dbObject = (DBObject) JSON.parse((String) streamsDatum.getDocument()); } else { try { ObjectNode node = mapper.valueToTree(streamsDatum.getDocument()); dbObject = (DBObject) JSON.parse(node.toString()); } catch (Exception ex) { LOGGER.error("Unsupported type: " + streamsDatum.getDocument().getClass(), ex); } } return dbObject; } 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()); } }