Refactor image writer

Core/src/org/sleuthkit/autopsy/casemodule/AddImageTask.java

@@ -26,6 +26,7 @@ import org.sleuthkit.autopsy.corecomponentinterfaces.DataSourceProcessorCallback
 import org.sleuthkit.autopsy.corecomponentinterfaces.DataSourceProcessorCallback.DataSourceProcessorResult;
 import org.sleuthkit.autopsy.corecomponentinterfaces.DataSourceProcessorProgressMonitor;
 import org.sleuthkit.autopsy.coreutils.Logger;
+import org.sleuthkit.autopsy.datasourceprocessors.ImageWriter;
 import org.sleuthkit.datamodel.Content;
 import org.sleuthkit.datamodel.Image;
 import org.sleuthkit.datamodel.SleuthkitJNI;
@@ -207,7 +208,9 @@ class AddImageTask implements Runnable {
                             errorMessages.add(verificationError);
                         }
                         if(! imageWriterPath.isEmpty()){
-                            Case.getCurrentCase().scheduleImageWriterFinish(imageId);
+                            // The ImageWriter object registers itself as an event listener and will 
+                            // stick around after this task is complete.
+                            ImageWriter writer = new ImageWriter(imageId);
                         }
                         newDataSources.add(newImage);
                     } else {

Core/src/org/sleuthkit/autopsy/casemodule/Case.java

@@ -287,7 +287,6 @@ public class Case implements SleuthkitCase.ErrorObserver {
     private static final int MIN_SECS_BETWEEN_TSK_ERROR_REPORTS = 60;
     private static final Logger logger = Logger.getLogger(Case.class.getName());
     private static final AutopsyEventPublisher eventPublisher = new AutopsyEventPublisher();
-    private final ImageWriter imageWriter = new ImageWriter();
     private static String appName;
     private static Case currentCase;
     private static CoordinationService.Lock currentCaseLock;
@@ -1439,7 +1438,6 @@ public class Case implements SleuthkitCase.ErrorObserver {
                 WindowManager.getDefault().getMainWindow().setCursor(Cursor.getPredefinedCursor(Cursor.WAIT_CURSOR));
             });
             IngestManager.getInstance().cancelAllIngestJobs(IngestJob.CancellationReason.CASE_CLOSED);
-            oldCase.closeImageWriter();
             completeCaseChange(null); //closes windows, etc   
             if (null != oldCase.tskErrorReporter) {
                 oldCase.tskErrorReporter.shutdown(); // stop listening for TSK errors for the old case
@@ -1616,22 +1614,6 @@ public class Case implements SleuthkitCase.ErrorObserver {
         }
         return currentCaseExecutor;
     }
-    
-    /**
-     * Register the ID of an image that is being copied using ImageWriter.
-     * This will cause the image to be finished after ingest is complete.
-     * @param imageID The image ID
-     */
-    void scheduleImageWriterFinish(long imageID){
-        imageWriter.addDataSourceId(imageID);
-    }
-    
-    /**
-     * Cancel all tasks associated with Image Writer
-     */
-    void closeImageWriter(){
-        imageWriter.close();
-    }
 
     /**
      * Gets the time zone(s) of the image data source(s) in this case.

Core/src/org/sleuthkit/autopsy/datasourceprocessors/ImageWriter.java

@@ -21,7 +21,6 @@ package org.sleuthkit.autopsy.datasourceprocessors;
 import com.google.common.util.concurrent.ThreadFactoryBuilder;
 import java.beans.PropertyChangeEvent;
 import java.beans.PropertyChangeListener;
-import java.util.HashSet;
 import java.util.concurrent.Executors;
 import java.util.concurrent.Future;
 import java.util.concurrent.ScheduledFuture;
@@ -41,192 +40,198 @@ import org.sleuthkit.datamodel.TskCoreException;
 
 /**
  * The ImageWriter class is used to complete VHD copies created from local disks
- * after the ingest process completes. 
+ * after the ingest process completes. The AddImageTask for this data source must have included 
+ * a non-empty imageWriterPath parameter to enable Image Writer.
  */
-public class ImageWriter {
+public class ImageWriter implements PropertyChangeListener{
     
     private final Logger logger = Logger.getLogger(ImageWriter.class.getName());
     
-    private final HashSet<Long> dataSourceIds = new HashSet<>();
+    private final Long dataSourceId;
-    private final Object dataSourceIdsLock;  // Get this lock before accessing dataSourceIds
+    private Long imageHandle = null;
     
-    private final HashSet<ScheduledFuture<?>> progressUpdaters = new HashSet<>();
+    private Future<?> finishTask;
-    private final HashSet<Long> imagesBeingFinished = new HashSet<>(); 
+    ProgressHandle progressHandle = null;
-    private final HashSet<ProgressHandle> progressBars = new HashSet<>();
+    ScheduledFuture<?> progressUpdateTask = null;
-    private final HashSet<Future<?>> finishTasksInProgress = new HashSet<>();
     private boolean isCancelled;
-    private final Object currentTasksLock; // Get this lock before accessing imagesBeingFinished, progressBars, progressUpdaters, finishTasksInProgress or isCancelled
+    private boolean isStarted;
+    private boolean isFinished;
+    private final Object currentTasksLock; // Get this lock before accessing finishTask, progressHandle, progressUpdateTask, isCancelled,
+                                           // isStarted, or isFinished
     
-    private boolean listenerStarted;
     private ScheduledThreadPoolExecutor periodicTasksExecutor = null;
     private final boolean doUI;
     
-    public ImageWriter(){
+    private static int numFinishJobsInProgress = 0;
-        dataSourceIdsLock = new Object();
+    
+    public ImageWriter(Long dataSourceId){
+        this.dataSourceId = dataSourceId;
+        
         currentTasksLock = new Object();
-        listenerStarted = false; 
         isCancelled = false;
+        isStarted = false;
+        isFinished = false;
+        progressHandle = null;
+        progressUpdateTask = null;
+        finishTask = null;
         
         doUI = RuntimeProperties.coreComponentsAreActive();
         if(doUI){
-            periodicTasksExecutor = new ScheduledThreadPoolExecutor(5, new ThreadFactoryBuilder().setNameFormat("image-writer-progress-update-%d").build()); //NON-NLS
+            periodicTasksExecutor = new ScheduledThreadPoolExecutor(1, new ThreadFactoryBuilder().setNameFormat("image-writer-progress-update-%d").build()); //NON-NLS
         }
+        
+        IngestManager.getInstance().addIngestJobEventListener(this);
+        Case.addEventSubscriber(Case.Events.CURRENT_CASE.toString(), this);
     }
     
     /**
-     * Creates a listener on IngestJobEvents if it hasn't already been started. 
+     * Handle the events:
-     * When a DataSourceAnalysisCompletedEvent arrives, if it matches
+     * DATA_SOURCE_ANALYSIS_COMPLETED - start the finish image process and clean up after it is complete
-     * the data source ID of an image that is using Image Writer, then finish the image
+     * CURRENT_CASE (case closing) - cancel the finish image process (if necessary)
-     * (fill in any gaps). The AddImageTask for this data source must have included 
-     * a non-empty imageWriterPath parameter to enable Image Writer.
      */
-    private synchronized void startListener(){
+    @Override
-        if(! listenerStarted){
+    public void propertyChange(PropertyChangeEvent evt) {
-            IngestManager.getInstance().addIngestJobEventListener(new PropertyChangeListener() {
+        if(evt.getPropertyName().equals(IngestManager.IngestJobEvent.DATA_SOURCE_ANALYSIS_COMPLETED.toString())){
-                @Override
+                        
-                public void propertyChange(PropertyChangeEvent evt) {
+            DataSourceAnalysisCompletedEvent event = (DataSourceAnalysisCompletedEvent)evt;
-                    if(evt.getPropertyName().equals(IngestManager.IngestJobEvent.DATA_SOURCE_ANALYSIS_COMPLETED.toString())){
-                        DataSourceAnalysisCompletedEvent event = (DataSourceAnalysisCompletedEvent)evt;
 
-                        if(event.getDataSource() != null){
+            if(event.getDataSource() != null){
-                            long imageId = event.getDataSource().getId();
+                long imageId = event.getDataSource().getId();
-                            String name = event.getDataSource().getName();
+                String name = event.getDataSource().getName();
-
+                
-                            // Check whether we need to run finishImage for this data source
+                // Check that the event corresponds to this datasource
-                            synchronized(dataSourceIdsLock){
+                if(imageId != dataSourceId){
-                                if( ! ImageWriter.this.dataSourceIds.contains(imageId)){
+                    return;
-                                    // Image writer was not used on this data source or we've already finished it
-                                    return;
-                                } else {
-                                    // Remove the imageId from the list here so we can't get past this point twice
-                                    // for the same image. Multiple DataSourceAnalysisCompletedEvent events can come from
-                                    // the same image if more ingest modules are run later, but the imageId is only added
-                                    // to the list during the intial task to add the image to the database.
-                                    ImageWriter.this.dataSourceIds.remove(imageId);
-                                }
-                            }
-                            logger.log(Level.INFO, String.format("Finishing VHD image for %s", 
-                                    event.getDataSource().getName())); //NON-NLS
-
-                            new Thread(() -> {
-                                try{
-                                    Image image = Case.getCurrentCase().getSleuthkitCase().getImageById(imageId);
-                                    ProgressHandle progressHandle = null;
-                                    ScheduledFuture<?> progressUpdateTask = null;
-                                    
-                                    if(doUI){
-                                        progressHandle = ProgressHandle.createHandle("Image writer - " + name);
-                                        progressHandle.start(100);
-                                        progressUpdateTask = periodicTasksExecutor.scheduleAtFixedRate(
-                                                new ProgressUpdateTask(progressHandle, image.getImageHandle()), 0, 250, TimeUnit.MILLISECONDS);
-                                    }
-
-                                    synchronized(currentTasksLock){
-                                        ImageWriter.this.imagesBeingFinished.add(image.getImageHandle());
-                                        
-                                        if(doUI){
-                                            if(isCancelled){
-                                                progressUpdateTask.cancel(true);
-                                                return;
-                                            }
-                                            ImageWriter.this.progressUpdaters.add(progressUpdateTask);
-                                            ImageWriter.this.progressBars.add(progressHandle);
-                                        }
-                                    }
-
-                                    // The added complexity here with the Future is because we absolutely need to make sure
-                                    // the call to finishImageWriter returns before allowing the TSK data structures to be freed
-                                    // during case close.
-                                    Future<?> finishTask = Executors.newSingleThreadExecutor().submit(() -> {
-                                        try{
-                                            SleuthkitJNI.finishImageWriter(image.getImageHandle());
-                                        } catch (TskCoreException ex){
-                                            logger.log(Level.SEVERE, "Error finishing VHD image", ex); //NON-NLS
-                                        }
-                                    });
-                                    synchronized(currentTasksLock){
-                                        ImageWriter.this.finishTasksInProgress.add(finishTask);
-                                    }
-                                        
-                                    // Wait for finishImageWriter to complete
-                                    try{
-                                        // The call to get() will happen twice if the user closes the case, which is ok
-                                        finishTask.get();
-                                    } catch (InterruptedException | ExecutionException ex){
-                                        logger.log(Level.SEVERE, "Error finishing VHD image", ex); //NON-NLS
-                                    }
-
-                                    synchronized(currentTasksLock){
-                                        ImageWriter.this.finishTasksInProgress.remove(finishTask);
-                                        ImageWriter.this.imagesBeingFinished.remove(image.getImageHandle());
-                                        
-                                        if(doUI){
-                                            progressUpdateTask.cancel(true);
-                                            ImageWriter.this.progressUpdaters.remove(progressUpdateTask);
-                                            progressHandle.finish();
-                                            ImageWriter.this.progressBars.remove(progressHandle);
-                                        }
-                                    }
-                                    
-                                    logger.log(Level.INFO, String.format("Finished writing VHD image for %s", event.getDataSource().getName())); //NON-NLS
-                                } catch (TskCoreException ex){
-                                    logger.log(Level.SEVERE, "Error finishing VHD image", ex); //NON-NLS
-                                }
-                            }).start();
-                        } else {
-                            logger.log(Level.SEVERE, "DataSourceAnalysisCompletedEvent did not contain a dataSource object"); //NON-NLS
-                        }
-                    }
                 }
-            });
+                new Thread(() -> {
+                    startFinishImage(name);
+                }).start();
+
+            } else {
+                logger.log(Level.SEVERE, "DataSourceAnalysisCompletedEvent did not contain a dataSource object"); //NON-NLS
+            }
+        }
+        else if(evt.getPropertyName().equals(Case.Events.CURRENT_CASE.toString())){
+            close();
         }
-        listenerStarted = true;
     }
     
-    /**
+    private void startFinishImage(String dataSourceName){
-     * Add a data source ID to the list of images to run finishImage on.
+        synchronized(currentTasksLock){
-     * Also starts the listener if needed.
+            // If we've already started the finish process for this datasource, return.
-     * @param id The dataSource/Image ID
+            // Multiple DataSourceAnalysisCompletedEvent events can come from
-     */
+            // the same image if more ingest modules are run later
-    public void addDataSourceId(Long id){
+            if(isStarted){
-        startListener();
+                return;
-        synchronized(dataSourceIdsLock){
+            } else {
-            dataSourceIds.add(id);
+                isStarted = true;
+            }
+        }
+
+        logger.log(Level.INFO, String.format("Finishing VHD image for %s", 
+                dataSourceName)); //NON-NLS
+
+        try{
+            Image image = Case.getCurrentCase().getSleuthkitCase().getImageById(dataSourceId);
+            imageHandle = image.getImageHandle();
+
+            synchronized(currentTasksLock){
+                if(isCancelled){
+                    return;
+                }
+
+                if(doUI){
+                    progressHandle = ProgressHandle.createHandle("Image writer - " + dataSourceName);
+                    progressHandle.start(100);
+                    progressUpdateTask = periodicTasksExecutor.scheduleAtFixedRate(
+                            new ProgressUpdateTask(progressHandle, image.getImageHandle()), 0, 250, TimeUnit.MILLISECONDS);
+                }
+
+                // The added complexity here with the Future is because we absolutely need to make sure
+                // the call to finishImageWriter returns before allowing the TSK data structures to be freed
+                // during case close.
+                numFinishJobsInProgress++;
+                finishTask = Executors.newSingleThreadExecutor().submit(() -> {
+                    try{
+                        SleuthkitJNI.finishImageWriter(imageHandle);
+                    } catch (TskCoreException ex){
+                        logger.log(Level.SEVERE, "Error finishing VHD image", ex); //NON-NLS
+                    }
+                });
+            }
+
+            // Wait for finishImageWriter to complete
+            try{
+                // The call to get() will happen twice if the user closes the case, which is ok
+                finishTask.get();
+            } catch (InterruptedException | ExecutionException ex){
+                logger.log(Level.SEVERE, "Error finishing VHD image", ex); //NON-NLS
+            }
+            numFinishJobsInProgress--;
+            
+            IngestManager.getInstance().removeIngestJobEventListener(this);
+            Case.removeEventSubscriber(Case.Events.CURRENT_CASE.toString(), this);
+            synchronized(currentTasksLock){
+                if(doUI && ! isCancelled){
+                    progressUpdateTask.cancel(true);
+                    progressHandle.finish();
+                }
+                isFinished = true;
+            }
+            
+
+            logger.log(Level.INFO, String.format("Finished writing VHD image for %s", dataSourceName)); //NON-NLS
+        } catch (TskCoreException | IllegalStateException ex){
+            logger.log(Level.SEVERE, "Error finishing VHD image", ex); //NON-NLS
         }
     }
     
     /**
-     * Stop any open progress update task, finish the progress bars, and tell
+     * Tell the finishImage process to stop and wait for it to do so.
-     * the finishImage process to stop
      */
-    public void close(){
+    private void close(){
         synchronized(currentTasksLock){
             isCancelled = true;
             
-            for(ScheduledFuture<?> task:ImageWriter.this.progressUpdaters){
+            if(imageHandle == null){
-                task.cancel(true);
+                // The case got closed during ingest (before the finish process could start)
+                return;
             }
             
-            for(Long handle:imagesBeingFinished){
+            if(!isFinished){
-                SleuthkitJNI.cancelFinishImage(handle);
+                SleuthkitJNI.cancelFinishImage(imageHandle);
                 logger.log(Level.SEVERE, "Case closed before VHD image could be finished"); //NON-NLS
-            }
             
-            // Wait for all the finish tasks to end
+                // Wait for the finish task to end
-            for(Future<?> task:ImageWriter.this.finishTasksInProgress){
                 try{
-                    task.get();
+                    finishTask.get();
                 } catch (InterruptedException | ExecutionException ex){
                     logger.log(Level.SEVERE, "Error finishing VHD image", ex); //NON-NLS
                 }
-            }
             
-            for(ProgressHandle progressHandle:ImageWriter.this.progressBars){
+                // Stop the progress bar and progress bar update task.
-                progressHandle.finish();
+                // The thread from startFinishImage will also stop these
+                // once the task completes, but we have to make absolutely sure
+                // this gets done before the Sleuthkit data structures are freed.
+                if(progressUpdateTask != null){
+                    progressUpdateTask.cancel(true);
+                }
+            
+                if(progressHandle != null){
+                    progressHandle.finish();
+                }
             }
         }
     }
     
+    /**
+     * Get the number of images currently being finished.
+     * @return number of images in progress
+     */
+    public static int numberOfJobsInProgress(){
+        return numFinishJobsInProgress;
+    }
+    
     /**
      * Task to query the Sleuthkit processing to get the percentage done. 
      */