bundles/org.simantics.utils.datastructures/src/org/simantics/utils/datastructures/file/DirectorySizeTracker.java

   1 /*******************************************************************************
   2  * Copyright (c) 2017 Association for Decentralized Information Management
   3  * in Industry THTH ry.
   4  * All rights reserved. This program and the accompanying materials
   5  * are made available under the terms of the Eclipse Public License v1.0
   6  * which accompanies this distribution, and is available at
   7  * http://www.eclipse.org/legal/epl-v10.html
   8  *
   9  * Contributors:
  10  *     Semantum Oy - #7330 - initial API and implementation
  11  *     Semantum Oy - #7551 - improved robustness
  12  *******************************************************************************/
  13 package org.simantics.utils.datastructures.file;
  14
  15 import static java.nio.file.StandardWatchEventKinds.ENTRY_CREATE;
  16 import static java.nio.file.StandardWatchEventKinds.ENTRY_DELETE;
  17 import static java.nio.file.StandardWatchEventKinds.ENTRY_MODIFY;
  18 import static java.nio.file.StandardWatchEventKinds.OVERFLOW;
  19
  20 import java.io.Closeable;
  21 import java.io.IOException;
  22 import java.nio.file.AccessDeniedException;
  23 import java.nio.file.FileSystems;
  24 import java.nio.file.FileVisitResult;
  25 import java.nio.file.FileVisitor;
  26 import java.nio.file.Files;
  27 import java.nio.file.NoSuchFileException;
  28 import java.nio.file.Path;
  29 import java.nio.file.SimpleFileVisitor;
  30 import java.nio.file.WatchEvent;
  31 import java.nio.file.WatchEvent.Kind;
  32 import java.nio.file.WatchKey;
  33 import java.nio.file.WatchService;
  34 import java.nio.file.attribute.BasicFileAttributes;
  35 import java.util.Set;
  36 import java.util.concurrent.atomic.AtomicInteger;
  37 import java.util.function.LongConsumer;
  38
  39 import org.simantics.databoard.util.BijectionMap;
  40 import org.simantics.utils.FileUtils;
  41 import org.simantics.utils.datastructures.MapSet;
  42 import org.slf4j.Logger;
  43 import org.slf4j.LoggerFactory;
  44
  45 import gnu.trove.map.TObjectLongMap;
  46 import gnu.trove.map.hash.TObjectLongHashMap;
  47
  48 /**
  49  * @author Tuukka Lehtonen
  50  * @since 1.30.0
  51  */
  52 public class DirectorySizeTracker implements Runnable, Closeable {
  53
  54         private static final Logger LOGGER = LoggerFactory.getLogger(DirectorySizeTracker.class);
  55
  56         private static final int TRACE_NONE = 0;
  57         @SuppressWarnings("unused")
  58         private static final int TRACE_DIRS = 1;
  59         @SuppressWarnings("unused")
  60         private static final int TRACE_EVENTS = 2;
  61         @SuppressWarnings("unused")
  62         private static final int TRACE_FILES = 3;
  63
  64         private static final Kind<?>[] ALL_EVENTS = { ENTRY_CREATE, ENTRY_DELETE, ENTRY_MODIFY };
  65
  66         private static AtomicInteger threadCounter = new AtomicInteger(1);
  67
  68         private LongConsumer consumer;
  69         private Kind<?>[] events;
  70
  71         private WatchService watcher;
  72
  73         private Object lock = new Object();
  74
  75         // For watched directory tracking
  76         private BijectionMap<WatchKey, Path> trackedDirs = new BijectionMap<>();
  77         private MapSet<Path, Path> subdirs = new MapSet.Hash<Path, Path>();
  78         private MapSet<Path, Path> files = new MapSet.Hash<Path, Path>();
  79
  80         // For directory and file size tracking
  81         private TObjectLongMap<Path> entrySizes = new TObjectLongHashMap<>(1024, 0.5f, -1L);
  82         private TObjectLongMap<Path> dirSizes = new TObjectLongHashMap<>(512, 0.5f, -1L);
  83         private long totalSize = 0L;
  84
  85         private int traceLevel = TRACE_NONE;
  86         private boolean running = true;
  87         private Thread thread;
  88
  89         public static DirectorySizeTracker startTracker(LongConsumer sizeChangeListener) throws IOException {
  90                 DirectorySizeTracker watcher = new DirectorySizeTracker(sizeChangeListener, ALL_EVENTS);
  91                 watcher.thread = new Thread(watcher, DirectorySizeTracker.class.getSimpleName() + "-" + threadCounter.get());
  92                 watcher.thread.start();
  93                 return watcher;
  94         }
  95
  96         private DirectorySizeTracker(LongConsumer sizeChangeListener, Kind<?>[] events) throws IOException {
  97                 this.consumer = sizeChangeListener;
  98                 this.events = events;
  99                 this.watcher = FileSystems.getDefault().newWatchService();
 100         }
 101
 102         public void close(boolean joinThread) throws InterruptedException {
 103                 running = false;
 104                 thread.interrupt();
 105                 if (joinThread)
 106                         thread.join();
 107         }
 108
 109         @Override
 110         public void close() throws IOException {
 111                 try {
 112                         close(true);
 113                 } catch (InterruptedException e) {
 114                         throw new IOException(e);
 115                 }
 116         }
 117
 118         /**
 119          * Register the given directory with the WatchService to listen for the
 120          * default set of events provided to the constructor.
 121          */
 122         public void track(Path dir) throws IOException {
 123                 synchronized (lock) {
 124                         if (trackedDirs.containsRight(dir))
 125                                 return;
 126                         if (traceLevel > 0)
 127                                 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);
 128                         Files.walkFileTree(dir, REGISTER);
 129                         if (traceLevel > 0)
 130                                 LOGGER.info("Now tracking entire directory " + dir + " with a total of " + entrySizes.size() + " files and " + trackedDirs.size() + " dirs with a total size of " + totalSize);
 131                 }
 132         }
 133
 134         private FileVisitor<Path> REGISTER = new SimpleFileVisitor<Path>() {
 135                 private Path currentDir;
 136
 137                 public java.nio.file.FileVisitResult preVisitDirectory(Path dir, BasicFileAttributes attrs) throws IOException {
 138                         currentDir = dir;
 139                         if (traceLevel > 2)
 140                                 LOGGER.info("Set current dir to " + currentDir);
 141                         return FileVisitResult.CONTINUE;
 142                 }
 143
 144                 public java.nio.file.FileVisitResult visitFile(Path file, BasicFileAttributes attrs) throws IOException {
 145                         registerFile(currentDir, file, attrs);
 146                         return FileVisitResult.CONTINUE;
 147                 }
 148
 149                 public java.nio.file.FileVisitResult visitFileFailed(Path file, IOException exc) throws IOException {
 150                         LOGGER.warn("Failed to visit file " + file, exc);
 151                         return FileVisitResult.CONTINUE;
 152                 }
 153
 154                 public java.nio.file.FileVisitResult postVisitDirectory(Path dir, IOException exc) throws IOException {
 155                         // Start tracking directories only after registering all entries
 156                         // within them.
 157                         Path parent = dir.getParent();
 158                         if (trackDir(dir))
 159                                 subdirs.add(parent, dir);
 160                         currentDir = parent;
 161                         return FileVisitResult.CONTINUE;
 162                 }
 163         };
 164
 165         /**
 166          * Register the given directory with the WatchService to listen for the
 167          * specified set of events.
 168          */
 169         private boolean trackDir(Path dir) throws IOException {
 170                 if (trackedDirs.containsRight(dir))
 171                         return false;
 172
 173                 WatchKey key = dir.register(watcher, events);
 174
 175                 if (traceLevel > 1) {
 176                         Path prev = trackedDirs.getRight(key);
 177                         if (prev == null) {
 178                                 LOGGER.info("Tracking new directory {}\n", dir);
 179                         } else {
 180                                 if (!dir.equals(prev)) {
 181                                         LOGGER.info("Tracked directory update: {} -> {}\n", prev, dir);
 182                                 }
 183                         }
 184                 }
 185
 186                 trackedDirs.map(key, dir);
 187                 return true;
 188         }
 189
 190         private boolean registerFile(Path dir, Path file, BasicFileAttributes attrs) {
 191                 if (files.add(dir, file)) {
 192                         long entrySize = attrs.size();
 193                         entrySizes.put(file, entrySize);
 194                         dirSizes.adjustOrPutValue(dir, entrySize, entrySize);
 195                         totalSize += entrySize;
 196                         if (traceLevel > 2)
 197                                 LOGGER.info("Registered file " + file + " size " + entrySize + " for total size " + totalSize);
 198                         return true;
 199                 } else {
 200                         long size = attrs.size();
 201                         long oldSize = entrySizes.put(file, size);
 202                         long sizeDelta = oldSize >= 0 ? size - oldSize : size;
 203                         totalSize += sizeDelta;
 204                         dirSizes.adjustOrPutValue(dir, sizeDelta, sizeDelta);
 205                         if (traceLevel > 2)
 206                                 LOGGER.info("Modified " + file + " size from " + oldSize + " to " + size
 207                                                 + " with delta " + sizeDelta + ", dir size = " + dirSizes.get(dir)
 208                                                 + ", total size = " + totalSize);
 209                 }
 210                 return false;
 211         }
 212
 213         private boolean unregisterFile(Path dir, Path file) {
 214                 long fileSize = entrySizes.remove(file);
 215                 if (fileSize >= 0) {
 216                         totalSize -= fileSize;
 217                         if (files.remove(dir, file))
 218                                 dirSizes.adjustValue(dir, -fileSize);
 219                         if (traceLevel > 2)
 220                                 LOGGER.info("Unregistered file " + file + " of size " + fileSize + ", dirSize = " + dirSizes.get(dir) + ", totalSize = " + totalSize);
 221                         return true;
 222                 }
 223                 return false;
 224         }
 225
 226         public void untrack(Path dir) {
 227                 synchronized (lock) {
 228                         if (!trackedDirs.containsRight(dir))
 229                                 return;
 230                         if (traceLevel > 0)
 231                                 LOGGER.info("Starting to untrack entire directory " + dir + " with total tracked size " + totalSize);
 232                         untrackTree(dir);
 233                         subdirs.remove(dir.getParent(), dir);
 234                         if (traceLevel > 0)
 235                                 LOGGER.info("Done untracking entire directory " + dir + " with total tracked size " + totalSize);
 236                 }
 237         }
 238
 239         private void untrackTree(Path dir) {
 240                 Set<Path> subdirs = this.subdirs.removeValues(dir);
 241                 for (Path subdir : subdirs)
 242                         untrackTree(subdir);
 243                 untrackDir(dir);
 244         }
 245
 246         private void untrackDir(Path dir) {
 247                 if (traceLevel > 1)
 248                         LOGGER.info("Untrack directory " + dir + " with total tracked size " + totalSize);
 249                 WatchKey key = trackedDirs.removeWithRight(dir);
 250                 if (key != null)
 251                         key.cancel();
 252
 253                 Set<Path> registeredFiles = files.removeValues(dir);
 254                 for (Path file : registeredFiles)
 255                         unregisterFile(dir, file);
 256                 dirSizes.remove(dir);
 257         }
 258
 259         @SuppressWarnings("unchecked")
 260         static <T> WatchEvent<T> cast(WatchEvent<?> event) {
 261                 return (WatchEvent<T>) event;
 262         }
 263
 264         /**
 265          * Process all events for keys queued to the watcher
 266          */
 267         void processEvents() {
 268                 while (running) {
 269                         // wait for key to be signaled
 270                         WatchKey key;
 271                         try {
 272                                 key = watcher.take();
 273                         } catch (InterruptedException x) {
 274                                 return;
 275                         }
 276
 277                         Path dir = trackedDirs.getRight(key);
 278                         if (dir == null) {
 279                                 LOGGER.error("WatchKey not registered: " + key);
 280                                 for (WatchEvent<?> event : key.pollEvents()) {
 281                                         WatchEvent.Kind<?> kind = event.kind();
 282                                         // TBD - provide example of how OVERFLOW event is handled
 283                                         if (kind == OVERFLOW)
 284                                                 continue;
 285                                         // Context for directory entry event is the file name of entry
 286                                         WatchEvent<Path> evt = cast(event);
 287                                         Path name = evt.context();
 288                                         LOGGER.error(String.format("MISSED EVENT: %s: %s", event.kind().name(), name));
 289                                 }
 290                                 boolean valid = key.reset();
 291                                 LOGGER.error("RESET KEY RETURNED: " + valid);
 292                                 continue;
 293                         }
 294
 295                         synchronized (lock) {
 296                                 for (WatchEvent<?> event : key.pollEvents()) {
 297                                         WatchEvent.Kind<?> kind = event.kind();
 298
 299                                         // TBD - provide example of how OVERFLOW event is handled
 300                                         if (kind == OVERFLOW)
 301                                                 continue;
 302
 303                                         // Context for directory entry event is the file name of entry
 304                                         WatchEvent<Path> evt = cast(event);
 305                                         Path name = evt.context();
 306                                         Path child = dir.resolve(name);
 307
 308                                         if (traceLevel > 1)
 309                                                 LOGGER.info(String.format("%s: %s", event.kind().name(), child));
 310
 311                                         if (kind == ENTRY_CREATE) {
 312                                                 try {
 313                                                         BasicFileAttributes attrs = Files.readAttributes(child, BasicFileAttributes.class);
 314                                                         if (attrs.isDirectory())
 315                                                                 track(child);
 316                                                         else if (attrs.isRegularFile()) {
 317                                                                 registerFile(dir, child, attrs);
 318                                                         }
 319                                                 } catch (NoSuchFileException ioe) {
 320                                                         // Intentionally ignore.
 321                                                         // The file was removed before we even had a chance to process it.
 322                                                 } catch (IOException ioe) {
 323                                                         LOGGER.error("Failed to read attribute for path " + child, ioe);
 324                                                 }
 325                                         } else if (kind == ENTRY_MODIFY) {
 326                                                 try {
 327                                                         BasicFileAttributes attrs = Files.readAttributes(child, BasicFileAttributes.class);
 328                                                         if (attrs.isRegularFile()) {
 329                                                                 registerFile(dir, child, attrs);
 330                                                         }
 331                                                 } catch (NoSuchFileException ioe) {
 332                                                         // It is possible that child is a directory that has been removed.
 333                                                         // In this case, just untrack the whole tree of child.
 334                                                         if (!entrySizes.containsKey(child)) {
 335                                                                 untrack(child);
 336                                                         }
 337                                                 } catch (AccessDeniedException e) {
 338                                                         // Ignore. This can happen e.g. with temporary locked Lucene index files.
 339                                                 } catch (IOException ioe) {
 340                                                         LOGGER.error("Failed to read attribute for path " + child, ioe);
 341                                                 }
 342                                         } else if (kind == ENTRY_DELETE) {
 343                                                 if (!unregisterFile(dir, child)) {
 344                                                         // This must have been a directory since it isn't registered as a file.
 345                                                         untrack(child);
 346                                                 }
 347                                         }
 348                                 }
 349
 350                                 // reset key and remove from set if directory no longer accessible
 351                                 boolean valid = key.reset();
 352                                 if (!valid) {
 353                                         if (traceLevel > 0)
 354                                                 LOGGER.info("WatchKey for dir " + dir + " is no longer valid. Untracking it.");
 355                                         untrack(dir);
 356                                         if (trackedDirs.isEmpty()) {
 357                                                 // all directories are inaccessible
 358                                                 break;
 359                                         }
 360                                 }
 361                         }
 362
 363                         if (traceLevel > 1)
 364                                 LOGGER.info("STATUS: Tracking a total of " + entrySizes.size() + " files and " + trackedDirs.size() + " dirs with a total size of " + totalSize);
 365                         if (consumer != null)
 366                                 consumer.accept(totalSize);
 367                 }
 368         }
 369
 370         @Override
 371         public void run() {
 372                 try {
 373                         while (running)
 374                                 processEvents();
 375                 } finally {
 376                         // It seems that for some reason cancelling all watch keys
 377                         // helps the watcher close itself. Otherwise closing the
 378                         // watch service seemed to get stuck easily.
 379                         cancelWatchKeys();
 380                         FileUtils.uncheckedClose(watcher);
 381                 }
 382         }
 383
 384         private void cancelWatchKeys() {
 385                 synchronized (lock) {
 386                         for (WatchKey key : trackedDirs.getLeftSet()) {
 387                                 key.cancel();
 388                                 key.pollEvents();
 389                         }
 390                 }
 391         }
 392
 393         /**
 394          * @return total size of the tracked directories in bytes
 395          */
 396         public long getTotalSize() {
 397                 return totalSize;
 398         }
 399
 400 //      public static void main(String[] args) {
 401 //              try (DirectorySizeTracker tracker = DirectorySizeTracker.startTracker(null)) {
 402 //                      tracker.track(Paths.get("d:/temp dir"));
 403 //                      LOGGER.info("AFTER TRACK: Total size from " + tracker.entrySizes.size() + " files and " + tracker.trackedDirs.size() + " directories is " + tracker.getTotalSize());
 404 //                      Thread.sleep(2000);
 405 //                      //tracker.untrack(Paths.get("d:/track-test"));
 406 //                      //LOGGER.info("AFTER UNTRACK: Total size from " + tracker.entrySizes.size() + " files and " + tracker.trackedDirs.size() + " directories is " + tracker.getTotalSize());
 407 //              } catch (IOException | InterruptedException e) {
 408 //                      LOGGER.error("test failed", e);
 409 //              }
 410 //      }
 411
 412 }