nifi/nifi-nar-bundles/nifi-provenance-repository-bundle/nifi-journaling-provenance-repository/src/main/java/org/apache/nifi/provenance/journaling/partition/QueuingPartitionManager.java - nifi - Git at Google

 /*
  * 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.nifi.provenance.journaling.partition;

 import java.io.File;
 import java.io.IOException;
 import java.util.ArrayList;
 import java.util.Collections;
 import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
 import java.util.concurrent.BlockingQueue;
 import java.util.concurrent.ExecutorService;
 import java.util.concurrent.LinkedBlockingQueue;
 import java.util.concurrent.ScheduledExecutorService;
 import java.util.concurrent.TimeUnit;
 import java.util.concurrent.atomic.AtomicLong;

 import org.apache.nifi.provenance.journaling.config.JournalingRepositoryConfig;
 import org.apache.nifi.provenance.journaling.index.IndexManager;
 import org.apache.nifi.util.Tuple;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 public class QueuingPartitionManager implements PartitionManager {

     private static final Logger logger = LoggerFactory.getLogger(QueuingPartitionManager.class);

     private final IndexManager indexManager;
     private final JournalingRepositoryConfig config;
     private final BlockingQueue<Partition> partitionQueue;
     private final JournalingPartition[] partitionArray;
     private final AtomicLong eventIdGenerator;
     private volatile boolean shutdown = false;

     private final Set<Partition> blackListedPartitions = Collections.synchronizedSet(new HashSet<Partition>());

     public QueuingPartitionManager(final IndexManager indexManager, final AtomicLong eventIdGenerator, final JournalingRepositoryConfig config, final ScheduledExecutorService workerExecutor, final ExecutorService compressionExecutor) throws IOException {
         this.indexManager = indexManager;
         this.config = config;
         this.eventIdGenerator = eventIdGenerator;

         // We can consider using a PriorityQueue here instead. Keep track of how many Partitions are being written
         // to for each container, as a container usually maps to a physical drive. Then, prioritize the queue
         // so that the partitions that belong to Container A get a higher priority than those belonging to Container B
         // if there are currently more partitions on Container B being written to (i.e., we prefer a partition for the
         // container that is the least used at this moment). Would require significant performance testing to see if it
         // really provides any benefit.
         this.partitionQueue = new LinkedBlockingQueue<>(config.getPartitionCount());
         this.partitionArray = new JournalingPartition[config.getPartitionCount()];

         final List<Tuple<String, File>> containerTuples = new ArrayList<>(config.getContainers().size());
         for ( final Map.Entry<String, File> entry : config.getContainers().entrySet() ) {
             containerTuples.add(new Tuple<>(entry.getKey(), entry.getValue()));
         }

         final Map<String, AtomicLong> containerSizes = new HashMap<>();
         for ( final String containerName : config.getContainers().keySet() ) {
             containerSizes.put(containerName, new AtomicLong(0L));
         }

         for (int i=0; i < config.getPartitionCount(); i++) {
             final Tuple<String, File> tuple = containerTuples.get(i % containerTuples.size());
             final File section = new File(tuple.getValue(), String.valueOf(i));

             final String containerName = tuple.getKey();
             final JournalingPartition partition = new JournalingPartition(indexManager, containerName, i,
                     section, config, containerSizes.get(containerName), compressionExecutor);
             partition.restore();
             partitionQueue.offer(partition);
             partitionArray[i] = partition;
         }

         workerExecutor.scheduleWithFixedDelay(new CheckBlackListedPartitions(), 30, 30, TimeUnit.SECONDS);
     }

     @Override
     public void shutdown() {
         this.shutdown = true;

         for ( final Partition partition : partitionArray ) {
             partition.shutdown();
         }
     }

     Partition nextPartition(final boolean writeAction, final boolean waitIfNeeded) {
         Partition partition = null;

         final List<Partition> partitionsSkipped = new ArrayList<>();
         try {
             while (partition == null) {
                 if (shutdown) {
                     throw new RuntimeException("Journaling Provenance Repository is shutting down");
                 }

                 try {
                     partition = partitionQueue.poll(1, TimeUnit.SECONDS);
                 } catch (final InterruptedException ie) {
                 }

                 if ( partition == null ) {
                     if ( blackListedPartitions.size() >= config.getPartitionCount() ) {
                         throw new RuntimeException("Cannot persist to the Journal Provenance Repository because all partitions have been blacklisted due to write failures");
                     }

                     // we are out of partitions. Add back all of the partitions that we skipped so we
                     // can try them again.
                     partitionQueue.addAll(partitionsSkipped);
                     partitionsSkipped.clear();
                 } else if (writeAction) {
                     if ( waitIfNeeded ) {
                         // determine if the container is full.
                         final String containerName = partition.getContainerName();
                         long desiredMaxContainerCapacity = config.getMaxCapacity(containerName);

                         // If no max capacity set for the container itself, use 1/N of repo max
                         // where N is the number of containers
                         if ( desiredMaxContainerCapacity == config.getMaxStorageCapacity() ) {
                             desiredMaxContainerCapacity = config.getMaxStorageCapacity() / config.getContainers().size();
                         }

                         // if the partition is more than 10% over its desired capacity, we don't want to write to it.
                         if ( partition.getContainerSize() > 1.1 * desiredMaxContainerCapacity ) {
                             partitionsSkipped.add(partition);
                             continue;
                         }
                     } else {
                         return null;
                     }
                 }
             }
         } finally {
             partitionQueue.addAll( partitionsSkipped );
         }

         return partition;
     }


     private void blackList(final Partition partition) {
         blackListedPartitions.add(partition);
     }

     @Override
     public <T> T withPartition(final PartitionAction<T> action, final boolean writeAction) throws IOException {
         final Partition partition = nextPartition(writeAction, true);

         boolean ioe = false;
         try {
             return action.perform(partition);
         } catch (final IOException e) {
             ioe = true;
             throw e;
         } finally {
             if ( ioe && writeAction ) {
                 blackList(partition);
             } else {
                 partitionQueue.offer(partition);
             }
         }
     }

     @Override
     public void withPartition(final VoidPartitionAction action, final boolean writeAction) throws IOException {
         final Partition partition = nextPartition(writeAction, true);

         boolean ioe = false;
         try {
             action.perform(partition);
         } catch (final IOException e) {
             ioe = true;
             throw e;
         } finally {
             if ( ioe && writeAction ) {
                 blackList(partition);
             } else {
                 partitionQueue.offer(partition);
             }
         }
     }


     @Override
     public <T> Set<T> withEachPartitionSerially(final PartitionAction<T> action, final boolean writeAction) throws IOException {
         if ( writeAction && blackListedPartitions.size() > 0 ) {
             throw new IOException("Cannot perform action {} because at least one partition has been blacklisted (i.e., writint to the partition failed)");
         }

         final Set<T> results = new HashSet<>(partitionArray.length);
         for ( final Partition partition : partitionArray ) {
             results.add( action.perform(partition) );
         }

         return results;
     }

     @Override
     public void withEachPartitionSerially(final VoidPartitionAction action, final boolean writeAction) throws IOException {
         if ( writeAction && blackListedPartitions.size() > 0 ) {
             throw new IOException("Cannot perform action {} because at least one partition has been blacklisted (i.e., writint to the partition failed)");
         }

         for ( final Partition partition : partitionArray ) {
             action.perform(partition);
         }
     }

     private long getTotalSize() {
         long totalSize = 0L;

         for ( final JournalingPartition partition : partitionArray ) {
             totalSize += partition.getPartitionSize();
         }

         for ( final String containerName : config.getContainers().keySet() ) {
             totalSize += indexManager.getSize(containerName);
         }

         return totalSize;
     }


     /**
      * Responsible for looking at partitions that have been marked as blacklisted and checking if they
      * are able to be written to now. If so, adds them back to the partition queue; otherwise, leaves
      * them as blacklisted
      */
     private class CheckBlackListedPartitions implements Runnable {
         @Override
         public void run() {
             final Set<Partition> reclaimed = new HashSet<>();

             final Set<Partition> partitions = new HashSet<>(blackListedPartitions);
             for ( final Partition partition : partitions ) {
                 final long nextId = eventIdGenerator.get();
                 if ( nextId <= 0 ) {
                     // we don't have an ID to use yet. Don't attempt to do anything yet.
                     return;
                 }

                 try {
                     partition.verifyWritable(nextId);
                     reclaimed.add(partition);
                  } catch (final IOException ioe) {
                      logger.debug("{} is still blackListed due to {}", partition, ioe);
                  }
             }

             // any partition that is reclaimable is now removed from the set of blacklisted
             // partitions and added back to our queue of partitions
             blackListedPartitions.removeAll(reclaimed);
             partitionQueue.addAll(reclaimed);
         }
     }


     @Override
     public void deleteEventsBasedOnSize() {
         final Map<String, List<JournalingPartition>> containerPartitionMap = new HashMap<>();

         for ( final JournalingPartition partition : partitionArray ) {
             final String container = partition.getContainerName();
             List<JournalingPartition> list = containerPartitionMap.get(container);
             if ( list == null ) {
                 list = new ArrayList<>();
                 containerPartitionMap.put(container, list);
             }

             list.add(partition);
         }

         int iterations = 0;
         for ( final String containerName : config.getContainers().keySet() ) {
             // continue as long as we need to delete data from this container.
             while (true) {
                 // don't hammer the disks if we can't delete anything
                 if ( iterations++ > 0 ) {
                     try {
                         Thread.sleep(1000L);
                     } catch (final InterruptedException ie) {}
                 }

                 final List<JournalingPartition> containerPartitions = containerPartitionMap.get(containerName);
                 final long containerSize = containerPartitions.get(0).getContainerSize();
                 final long maxContainerCapacity = config.getMaxCapacity(containerName);
                 if ( containerSize < maxContainerCapacity ) {
                     break;
                 }

                 logger.debug("Container {} exceeds max capacity of {} bytes with a size of {} bytes; deleting oldest events", containerName, maxContainerCapacity, containerSize);

                 // container is too large. Delete oldest journal from each partition in this container.
                 for ( final Partition partition : containerPartitions ) {
                     try {
                         partition.deleteOldest();
                     } catch (final IOException ioe) {
                         logger.error("Failed to delete events from {} due to {}", partition, ioe.toString());
                         if ( logger.isDebugEnabled() ) {
                             logger.error("", ioe);
                         }
                     }
                 }
             }
         }

         long totalSize;
         iterations = 0;
         while ((totalSize = getTotalSize()) >= config.getMaxStorageCapacity()) {
             logger.debug("Provenance Repository exceeds max capacity of {} bytes with a size of {}; deleting oldest events", config.getMaxStorageCapacity(), totalSize);

             // don't hammer the disks if we can't delete anything
             if ( iterations++ > 0 ) {
                 try {
                     Thread.sleep(1000L);
                 } catch (final InterruptedException ie) {}
             }

             for ( final Partition partition : partitionArray ) {
                 try {
                     partition.deleteOldest();
                 } catch (final IOException ioe) {
                     logger.error("Failed to delete events from {} due to {}", partition, ioe.toString());
                     if ( logger.isDebugEnabled() ) {
                         logger.error("", ioe);
                     }
                 }
             }

             // don't hammer the disks if we can't delete anything
             try {
                 Thread.sleep(1000L);
             } catch (final InterruptedException ie) {}
         }
     }
 }
	/*
	* 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.nifi.provenance.journaling.partition;

	import java.io.File;
	import java.io.IOException;
	import java.util.ArrayList;
	import java.util.Collections;
	import java.util.HashMap;
	import java.util.HashSet;
	import java.util.List;
	import java.util.Map;
	import java.util.Set;
	import java.util.concurrent.BlockingQueue;
	import java.util.concurrent.ExecutorService;
	import java.util.concurrent.LinkedBlockingQueue;
	import java.util.concurrent.ScheduledExecutorService;
	import java.util.concurrent.TimeUnit;
	import java.util.concurrent.atomic.AtomicLong;

	import org.apache.nifi.provenance.journaling.config.JournalingRepositoryConfig;
	import org.apache.nifi.provenance.journaling.index.IndexManager;
	import org.apache.nifi.util.Tuple;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	public class QueuingPartitionManager implements PartitionManager {

	private static final Logger logger = LoggerFactory.getLogger(QueuingPartitionManager.class);

	private final IndexManager indexManager;
	private final JournalingRepositoryConfig config;
	private final BlockingQueue<Partition> partitionQueue;
	private final JournalingPartition[] partitionArray;
	private final AtomicLong eventIdGenerator;
	private volatile boolean shutdown = false;

	private final Set<Partition> blackListedPartitions = Collections.synchronizedSet(new HashSet<Partition>());

	public QueuingPartitionManager(final IndexManager indexManager, final AtomicLong eventIdGenerator, final JournalingRepositoryConfig config, final ScheduledExecutorService workerExecutor, final ExecutorService compressionExecutor) throws IOException {
	this.indexManager = indexManager;
	this.config = config;
	this.eventIdGenerator = eventIdGenerator;

	// We can consider using a PriorityQueue here instead. Keep track of how many Partitions are being written
	// to for each container, as a container usually maps to a physical drive. Then, prioritize the queue
	// so that the partitions that belong to Container A get a higher priority than those belonging to Container B
	// if there are currently more partitions on Container B being written to (i.e., we prefer a partition for the
	// container that is the least used at this moment). Would require significant performance testing to see if it
	// really provides any benefit.
	this.partitionQueue = new LinkedBlockingQueue<>(config.getPartitionCount());
	this.partitionArray = new JournalingPartition[config.getPartitionCount()];

	final List<Tuple<String, File>> containerTuples = new ArrayList<>(config.getContainers().size());
	for ( final Map.Entry<String, File> entry : config.getContainers().entrySet() ) {
	containerTuples.add(new Tuple<>(entry.getKey(), entry.getValue()));
	}

	final Map<String, AtomicLong> containerSizes = new HashMap<>();
	for ( final String containerName : config.getContainers().keySet() ) {
	containerSizes.put(containerName, new AtomicLong(0L));
	}

	for (int i=0; i < config.getPartitionCount(); i++) {
	final Tuple<String, File> tuple = containerTuples.get(i % containerTuples.size());
	final File section = new File(tuple.getValue(), String.valueOf(i));

	final String containerName = tuple.getKey();
	final JournalingPartition partition = new JournalingPartition(indexManager, containerName, i,
	section, config, containerSizes.get(containerName), compressionExecutor);
	partition.restore();
	partitionQueue.offer(partition);
	partitionArray[i] = partition;
	}

	workerExecutor.scheduleWithFixedDelay(new CheckBlackListedPartitions(), 30, 30, TimeUnit.SECONDS);
	}

	@Override
	public void shutdown() {
	this.shutdown = true;

	for ( final Partition partition : partitionArray ) {
	partition.shutdown();
	}
	}

	Partition nextPartition(final boolean writeAction, final boolean waitIfNeeded) {
	Partition partition = null;

	final List<Partition> partitionsSkipped = new ArrayList<>();
	try {
	while (partition == null) {
	if (shutdown) {
	throw new RuntimeException("Journaling Provenance Repository is shutting down");
	}

	try {
	partition = partitionQueue.poll(1, TimeUnit.SECONDS);
	} catch (final InterruptedException ie) {
	}

	if ( partition == null ) {
	if ( blackListedPartitions.size() >= config.getPartitionCount() ) {
	throw new RuntimeException("Cannot persist to the Journal Provenance Repository because all partitions have been blacklisted due to write failures");
	}

	// we are out of partitions. Add back all of the partitions that we skipped so we
	// can try them again.
	partitionQueue.addAll(partitionsSkipped);
	partitionsSkipped.clear();
	} else if (writeAction) {
	if ( waitIfNeeded ) {
	// determine if the container is full.
	final String containerName = partition.getContainerName();
	long desiredMaxContainerCapacity = config.getMaxCapacity(containerName);

	// If no max capacity set for the container itself, use 1/N of repo max
	// where N is the number of containers
	if ( desiredMaxContainerCapacity == config.getMaxStorageCapacity() ) {
	desiredMaxContainerCapacity = config.getMaxStorageCapacity() / config.getContainers().size();
	}

	// if the partition is more than 10% over its desired capacity, we don't want to write to it.
	if ( partition.getContainerSize() > 1.1 * desiredMaxContainerCapacity ) {
	partitionsSkipped.add(partition);
	continue;
	}
	} else {
	return null;
	}
	}
	}
	} finally {
	partitionQueue.addAll( partitionsSkipped );
	}

	return partition;
	}


	private void blackList(final Partition partition) {
	blackListedPartitions.add(partition);
	}

	@Override
	public <T> T withPartition(final PartitionAction<T> action, final boolean writeAction) throws IOException {
	final Partition partition = nextPartition(writeAction, true);

	boolean ioe = false;
	try {
	return action.perform(partition);
	} catch (final IOException e) {
	ioe = true;
	throw e;
	} finally {
	if ( ioe && writeAction ) {
	blackList(partition);
	} else {
	partitionQueue.offer(partition);
	}
	}
	}

	@Override
	public void withPartition(final VoidPartitionAction action, final boolean writeAction) throws IOException {
	final Partition partition = nextPartition(writeAction, true);

	boolean ioe = false;
	try {
	action.perform(partition);
	} catch (final IOException e) {
	ioe = true;
	throw e;
	} finally {
	if ( ioe && writeAction ) {
	blackList(partition);
	} else {
	partitionQueue.offer(partition);
	}
	}
	}


	@Override
	public <T> Set<T> withEachPartitionSerially(final PartitionAction<T> action, final boolean writeAction) throws IOException {
	if ( writeAction && blackListedPartitions.size() > 0 ) {
	throw new IOException("Cannot perform action {} because at least one partition has been blacklisted (i.e., writint to the partition failed)");
	}

	final Set<T> results = new HashSet<>(partitionArray.length);
	for ( final Partition partition : partitionArray ) {
	results.add( action.perform(partition) );
	}

	return results;
	}

	@Override
	public void withEachPartitionSerially(final VoidPartitionAction action, final boolean writeAction) throws IOException {
	if ( writeAction && blackListedPartitions.size() > 0 ) {
	throw new IOException("Cannot perform action {} because at least one partition has been blacklisted (i.e., writint to the partition failed)");
	}

	for ( final Partition partition : partitionArray ) {
	action.perform(partition);
	}
	}

	private long getTotalSize() {
	long totalSize = 0L;

	for ( final JournalingPartition partition : partitionArray ) {
	totalSize += partition.getPartitionSize();
	}

	for ( final String containerName : config.getContainers().keySet() ) {
	totalSize += indexManager.getSize(containerName);
	}

	return totalSize;
	}


	/**
	* Responsible for looking at partitions that have been marked as blacklisted and checking if they
	* are able to be written to now. If so, adds them back to the partition queue; otherwise, leaves
	* them as blacklisted
	*/
	private class CheckBlackListedPartitions implements Runnable {
	@Override
	public void run() {
	final Set<Partition> reclaimed = new HashSet<>();

	final Set<Partition> partitions = new HashSet<>(blackListedPartitions);
	for ( final Partition partition : partitions ) {
	final long nextId = eventIdGenerator.get();
	if ( nextId <= 0 ) {
	// we don't have an ID to use yet. Don't attempt to do anything yet.
	return;
	}

	try {
	partition.verifyWritable(nextId);
	reclaimed.add(partition);
	} catch (final IOException ioe) {
	logger.debug("{} is still blackListed due to {}", partition, ioe);
	}
	}

	// any partition that is reclaimable is now removed from the set of blacklisted
	// partitions and added back to our queue of partitions
	blackListedPartitions.removeAll(reclaimed);
	partitionQueue.addAll(reclaimed);
	}
	}


	@Override
	public void deleteEventsBasedOnSize() {
	final Map<String, List<JournalingPartition>> containerPartitionMap = new HashMap<>();

	for ( final JournalingPartition partition : partitionArray ) {
	final String container = partition.getContainerName();
	List<JournalingPartition> list = containerPartitionMap.get(container);
	if ( list == null ) {
	list = new ArrayList<>();
	containerPartitionMap.put(container, list);
	}

	list.add(partition);
	}

	int iterations = 0;
	for ( final String containerName : config.getContainers().keySet() ) {
	// continue as long as we need to delete data from this container.
	while (true) {
	// don't hammer the disks if we can't delete anything
	if ( iterations++ > 0 ) {
	try {
	Thread.sleep(1000L);
	} catch (final InterruptedException ie) {}
	}

	final List<JournalingPartition> containerPartitions = containerPartitionMap.get(containerName);
	final long containerSize = containerPartitions.get(0).getContainerSize();
	final long maxContainerCapacity = config.getMaxCapacity(containerName);
	if ( containerSize < maxContainerCapacity ) {
	break;
	}

	logger.debug("Container {} exceeds max capacity of {} bytes with a size of {} bytes; deleting oldest events", containerName, maxContainerCapacity, containerSize);

	// container is too large. Delete oldest journal from each partition in this container.
	for ( final Partition partition : containerPartitions ) {
	try {
	partition.deleteOldest();
	} catch (final IOException ioe) {
	logger.error("Failed to delete events from {} due to {}", partition, ioe.toString());
	if ( logger.isDebugEnabled() ) {
	logger.error("", ioe);
	}
	}
	}
	}
	}

	long totalSize;
	iterations = 0;
	while ((totalSize = getTotalSize()) >= config.getMaxStorageCapacity()) {
	logger.debug("Provenance Repository exceeds max capacity of {} bytes with a size of {}; deleting oldest events", config.getMaxStorageCapacity(), totalSize);

	// don't hammer the disks if we can't delete anything
	if ( iterations++ > 0 ) {
	try {
	Thread.sleep(1000L);
	} catch (final InterruptedException ie) {}
	}

	for ( final Partition partition : partitionArray ) {
	try {
	partition.deleteOldest();
	} catch (final IOException ioe) {
	logger.error("Failed to delete events from {} due to {}", partition, ioe.toString());
	if ( logger.isDebugEnabled() ) {
	logger.error("", ioe);
	}
	}
	}

	// don't hammer the disks if we can't delete anything
	try {
	Thread.sleep(1000L);
	} catch (final InterruptedException ie) {}
	}
	}
	}