--- /dev/null
+/*******************************************************************************
+ * Copyright (c) 2017 Association for Decentralized Information Management
+ * in Industry THTH ry.
+ * All rights reserved. This program and the accompanying materials
+ * are made available under the terms of the Eclipse Public License v1.0
+ * which accompanies this distribution, and is available at
+ * http://www.eclipse.org/legal/epl-v10.html
+ *
+ * Contributors:
+ * Semantum Oy - #7330 - initial API and implementation
+ *******************************************************************************/
+package org.simantics.utils.datastructures.file;
+
+import static java.nio.file.StandardWatchEventKinds.ENTRY_CREATE;
+import static java.nio.file.StandardWatchEventKinds.ENTRY_DELETE;
+import static java.nio.file.StandardWatchEventKinds.ENTRY_MODIFY;
+import static java.nio.file.StandardWatchEventKinds.OVERFLOW;
+
+import java.io.Closeable;
+import java.io.IOException;
+import java.nio.file.FileSystems;
+import java.nio.file.FileVisitResult;
+import java.nio.file.FileVisitor;
+import java.nio.file.Files;
+import java.nio.file.NoSuchFileException;
+import java.nio.file.Path;
+import java.nio.file.SimpleFileVisitor;
+import java.nio.file.WatchEvent;
+import java.nio.file.WatchEvent.Kind;
+import java.nio.file.WatchKey;
+import java.nio.file.WatchService;
+import java.nio.file.attribute.BasicFileAttributes;
+import java.util.Set;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.function.LongConsumer;
+
+import org.simantics.databoard.util.BijectionMap;
+import org.simantics.utils.FileUtils;
+import org.simantics.utils.datastructures.MapSet;
+import org.slf4j.Logger;
+import org.slf4j.LoggerFactory;
+
+import gnu.trove.map.TObjectLongMap;
+import gnu.trove.map.hash.TObjectLongHashMap;
+
+/**
+ * @author Tuukka Lehtonen
+ * @since 1.30.0
+ */
+public class DirectorySizeTracker implements Runnable, Closeable {
+
+ private static final Logger LOGGER = LoggerFactory.getLogger(DirectorySizeTracker.class);
+
+ private static final int TRACE_NONE = 0;
+ @SuppressWarnings("unused")
+ private static final int TRACE_DIRS = 1;
+ @SuppressWarnings("unused")
+ private static final int TRACE_EVENTS = 2;
+ @SuppressWarnings("unused")
+ private static final int TRACE_FILES = 3;
+
+ private static final Kind<?>[] ALL_EVENTS = { ENTRY_CREATE, ENTRY_DELETE, ENTRY_MODIFY };
+
+ private static AtomicInteger threadCounter = new AtomicInteger();
+
+ private LongConsumer consumer;
+ private Kind<?>[] events;
+
+ private WatchService watcher;
+
+ private Object lock = new Object();
+
+ // For watched directory tracking
+ private BijectionMap<WatchKey, Path> trackedDirs = new BijectionMap<>();
+ private MapSet<Path, Path> subdirs = new MapSet.Hash<Path, Path>();
+ private MapSet<Path, Path> files = new MapSet.Hash<Path, Path>();
+
+ // For directory and file size tracking
+ private TObjectLongMap<Path> entrySizes = new TObjectLongHashMap<>(1024, 0.5f, -1L);
+ private TObjectLongMap<Path> dirSizes = new TObjectLongHashMap<>(512, 0.5f, -1L);
+ private long totalSize = 0L;
+
+ private int traceLevel = TRACE_NONE;
+ private boolean running = true;
+ private Thread thread;
+
+ public static DirectorySizeTracker startTracker(LongConsumer sizeChangeListener) throws IOException {
+ DirectorySizeTracker watcher = new DirectorySizeTracker(sizeChangeListener, ALL_EVENTS);
+ watcher.thread = new Thread(watcher, DirectorySizeTracker.class.getSimpleName() + threadCounter.get());
+ watcher.thread.start();
+ return watcher;
+ }
+
+ private DirectorySizeTracker(LongConsumer sizeChangeListener, Kind<?>[] events) throws IOException {
+ this.consumer = sizeChangeListener;
+ this.events = events;
+ this.watcher = FileSystems.getDefault().newWatchService();
+ }
+
+ public void close(boolean joinThread) throws InterruptedException {
+ running = false;
+ thread.interrupt();
+ if (joinThread)
+ thread.join();
+ }
+
+ @Override
+ public void close() throws IOException {
+ try {
+ close(true);
+ } catch (InterruptedException e) {
+ throw new IOException(e);
+ }
+ }
+
+ /**
+ * Register the given directory with the WatchService to listen for the
+ * default set of events provided to the constructor.
+ */
+ public void track(Path dir) throws IOException {
+ synchronized (lock) {
+ if (trackedDirs.containsRight(dir))
+ return;
+ if (traceLevel > 0)
+ LOGGER.info("Starting to track entire directory " + dir + " with a total of " + entrySizes.size() + " files and " + trackedDirs.size() + " dirs with a total size of " + totalSize);
+ Files.walkFileTree(dir, REGISTER);
+ if (traceLevel > 0)
+ LOGGER.info("Now tracking entire directory " + dir + " with a total of " + entrySizes.size() + " files and " + trackedDirs.size() + " dirs with a total size of " + totalSize);
+ }
+ }
+
+ private FileVisitor<Path> REGISTER = new SimpleFileVisitor<Path>() {
+ private Path currentDir;
+
+ public java.nio.file.FileVisitResult preVisitDirectory(Path dir, BasicFileAttributes attrs) throws IOException {
+ currentDir = dir;
+ if (traceLevel > 2)
+ LOGGER.info("Set current dir to " + currentDir);
+ return FileVisitResult.CONTINUE;
+ }
+
+ public java.nio.file.FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
+ registerFile(currentDir, file, attrs);
+ return FileVisitResult.CONTINUE;
+ }
+
+ public java.nio.file.FileVisitResult visitFileFailed(Path file, IOException exc) throws IOException {
+ LOGGER.warn("Failed to visit file " + file, exc);
+ return FileVisitResult.CONTINUE;
+ }
+
+ public java.nio.file.FileVisitResult postVisitDirectory(Path dir, IOException exc) throws IOException {
+ // Start tracking directories only after registering all entries
+ // within them.
+ Path parent = dir.getParent();
+ if (trackDir(dir))
+ subdirs.add(parent, dir);
+ currentDir = parent;
+ return FileVisitResult.CONTINUE;
+ }
+ };
+
+ /**
+ * Register the given directory with the WatchService to listen for the
+ * specified set of events.
+ */
+ private boolean trackDir(Path dir) throws IOException {
+ if (trackedDirs.containsRight(dir))
+ return false;
+
+ WatchKey key = dir.register(watcher, events);
+
+ if (traceLevel > 1) {
+ Path prev = trackedDirs.getRight(key);
+ if (prev == null) {
+ LOGGER.info("Tracking new directory {}\n", dir);
+ } else {
+ if (!dir.equals(prev)) {
+ LOGGER.info("Tracked directory update: {} -> {}\n", prev, dir);
+ }
+ }
+ }
+
+ trackedDirs.map(key, dir);
+ return true;
+ }
+
+ private boolean registerFile(Path dir, Path file, BasicFileAttributes attrs) {
+ if (files.add(dir, file)) {
+ long entrySize = attrs.size();
+ entrySizes.put(file, entrySize);
+ dirSizes.adjustOrPutValue(dir, entrySize, entrySize);
+ totalSize += entrySize;
+ if (traceLevel > 2)
+ LOGGER.info("Registered file " + file + " size " + entrySize + " for total size " + totalSize);
+ return true;
+ } else {
+ long size = attrs.size();
+ long oldSize = entrySizes.put(file, size);
+ long sizeDelta = oldSize >= 0 ? size - oldSize : size;
+ totalSize += sizeDelta;
+ dirSizes.adjustOrPutValue(dir, sizeDelta, sizeDelta);
+ if (traceLevel > 2)
+ LOGGER.info("Modified " + file + " size from " + oldSize + " to " + size
+ + " with delta " + sizeDelta + ", dir size = " + dirSizes.get(dir)
+ + ", total size = " + totalSize);
+ }
+ return false;
+ }
+
+ private boolean unregisterFile(Path dir, Path file) {
+ long fileSize = entrySizes.remove(file);
+ if (fileSize >= 0) {
+ totalSize -= fileSize;
+ if (files.remove(dir, file))
+ dirSizes.adjustValue(dir, -fileSize);
+ if (traceLevel > 2)
+ LOGGER.info("Unregistered file " + file + " of size " + fileSize + ", dirSize = " + dirSizes.get(dir) + ", totalSize = " + totalSize);
+ return true;
+ }
+ return false;
+ }
+
+ public void untrack(Path dir) {
+ synchronized (lock) {
+ if (!trackedDirs.containsRight(dir))
+ return;
+ if (traceLevel > 0)
+ LOGGER.info("Starting to untrack entire directory " + dir + " with total tracked size " + totalSize);
+ untrackTree(dir);
+ subdirs.remove(dir.getParent(), dir);
+ if (traceLevel > 0)
+ LOGGER.info("Done untracking entire directory " + dir + " with total tracked size " + totalSize);
+ }
+ }
+
+ private void untrackTree(Path dir) {
+ Set<Path> subdirs = this.subdirs.removeValues(dir);
+ for (Path subdir : subdirs)
+ untrackTree(subdir);
+ untrackDir(dir);
+ }
+
+ private void untrackDir(Path dir) {
+ if (traceLevel > 1)
+ LOGGER.info("Untrack directory " + dir + " with total tracked size " + totalSize);
+ WatchKey key = trackedDirs.removeWithRight(dir);
+ if (key != null)
+ key.cancel();
+
+ Set<Path> registeredFiles = files.removeValues(dir);
+ for (Path file : registeredFiles)
+ unregisterFile(dir, file);
+ dirSizes.remove(dir);
+ }
+
+ @SuppressWarnings("unchecked")
+ static <T> WatchEvent<T> cast(WatchEvent<?> event) {
+ return (WatchEvent<T>) event;
+ }
+
+ /**
+ * Process all events for keys queued to the watcher
+ */
+ void processEvents() {
+ while (running) {
+ // wait for key to be signaled
+ WatchKey key;
+ try {
+ key = watcher.take();
+ } catch (InterruptedException x) {
+ return;
+ }
+
+ Path dir = trackedDirs.getRight(key);
+ if (dir == null) {
+ LOGGER.error("WatchKey not registered: " + key);
+ continue;
+ }
+
+ synchronized (lock) {
+ for (WatchEvent<?> event : key.pollEvents()) {
+ WatchEvent.Kind<?> kind = event.kind();
+
+ // TBD - provide example of how OVERFLOW event is handled
+ if (kind == OVERFLOW)
+ continue;
+
+ // Context for directory entry event is the file name of entry
+ WatchEvent<Path> evt = cast(event);
+ Path name = evt.context();
+ Path child = dir.resolve(name);
+
+ if (traceLevel > 1)
+ LOGGER.info(String.format("%s: %s", event.kind().name(), child));
+
+ if (kind == ENTRY_CREATE) {
+ try {
+ BasicFileAttributes attrs = Files.readAttributes(child, BasicFileAttributes.class);
+ if (attrs.isDirectory())
+ track(child);
+ else if (attrs.isRegularFile()) {
+ registerFile(dir, child, attrs);
+ }
+ } catch (IOException ioe) {
+ LOGGER.error("Failed to read attribute for path " + child, ioe);
+ }
+ } else if (kind == ENTRY_MODIFY) {
+ try {
+ BasicFileAttributes attrs = Files.readAttributes(child, BasicFileAttributes.class);
+ if (attrs.isRegularFile()) {
+ registerFile(dir, child, attrs);
+ }
+ } catch (NoSuchFileException ioe) {
+ // It is possible that child is a directory that has been removed.
+ // In this case, just untrack the whole tree of child.
+ if (!entrySizes.containsKey(child)) {
+ untrack(child);
+ }
+ } catch (IOException ioe) {
+ LOGGER.error("Failed to read attribute for path " + child, ioe);
+ }
+ } else if (kind == ENTRY_DELETE) {
+ if (!unregisterFile(dir, child)) {
+ // This must have been a directory since it isn't registered as a file.
+ untrack(child);
+ }
+ }
+ }
+
+ // reset key and remove from set if directory no longer accessible
+ boolean valid = key.reset();
+ if (!valid) {
+ if (traceLevel > 0)
+ LOGGER.info("WatchKey for dir " + dir + " is no longer valid. Untracking it.");
+ untrack(dir);
+ if (trackedDirs.isEmpty()) {
+ // all directories are inaccessible
+ break;
+ }
+ }
+ }
+
+ if (traceLevel > 1)
+ LOGGER.info("STATUS: Tracking a total of " + entrySizes.size() + " files and " + trackedDirs.size() + " dirs with a total size of " + totalSize);
+ if (consumer != null)
+ consumer.accept(totalSize);
+ }
+ }
+
+ @Override
+ public void run() {
+ try {
+ while (running)
+ processEvents();
+ } finally {
+ FileUtils.uncheckedClose(watcher);
+ }
+ }
+
+ /**
+ * @return total size of the tracked directories in bytes
+ */
+ public long getTotalSize() {
+ return totalSize;
+ }
+
+// public static void main(String[] args) {
+// try {
+// DirectorySizeTracker tracker = DirectorySizeTracker.startTracker(null);
+// tracker.track(Paths.get("d:/track-test"));
+// LOGGER.info("AFTER TRACK: Total size from " + tracker.entrySizes.size() + " files and " + tracker.trackedDirs.size() + " directories is " + tracker.getTotalSize());
+// Thread.sleep(2000);
+// //tracker.untrack(Paths.get("d:/track-test"));
+// //LOGGER.info("AFTER UNTRACK: Total size from " + tracker.entrySizes.size() + " files and " + tracker.trackedDirs.size() + " directories is " + tracker.getTotalSize());
+// } catch (IOException | InterruptedException e) {
+// LOGGER.error("test failed", e);
+// }
+// }
+
+}
\ No newline at end of file
Code Review / simantics / platform.git / commitdiff