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