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