CollaborationMonitor.java

Go to the documentation of this file.

/*
 * Autopsy Forensic Browser
 *
 * Copyright 2015-2021 Basis Technology Corp.
 * Contact: carrier <at> sleuthkit <dot> org
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.sleuthkit.autopsy.casemodule;

import com.google.common.util.concurrent.ThreadFactoryBuilder;
import java.beans.PropertyChangeEvent;
import java.beans.PropertyChangeListener;
import java.io.Serializable;
import java.time.Duration;
import java.time.Instant;
import java.util.Collections;
import java.util.EnumSet;
import java.util.HashMap;
import java.util.Iterator;
import java.util.Map;
import java.util.Set;
import java.util.UUID;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.logging.Level;
import org.netbeans.api.progress.ProgressHandle;
import org.openide.util.NbBundle;
import org.sleuthkit.autopsy.casemodule.events.AddingDataSourceEvent;
import org.sleuthkit.autopsy.casemodule.events.AddingDataSourceFailedEvent;
import org.sleuthkit.autopsy.casemodule.events.DataSourceAddedEvent;
import org.sleuthkit.autopsy.coreutils.Logger;
import org.sleuthkit.autopsy.coreutils.NetworkUtils;
import org.sleuthkit.autopsy.events.AutopsyEvent;
import org.sleuthkit.autopsy.events.AutopsyEventException;
import org.sleuthkit.autopsy.events.AutopsyEventPublisher;
import org.sleuthkit.autopsy.ingest.IngestManager;
import org.sleuthkit.autopsy.ingest.events.DataSourceAnalysisCompletedEvent;
import org.sleuthkit.autopsy.ingest.events.DataSourceAnalysisStartedEvent;
import org.sleuthkit.datamodel.Content;

final class CollaborationMonitor {

    private static final String EVENT_CHANNEL_NAME = "%s-Collaboration-Monitor-Events"; //NON-NLS
    private static final String COLLABORATION_MONITOR_EVENT = "COLLABORATION_MONITOR_EVENT"; //NON-NLS
    private static final Set<Case.Events> CASE_EVENTS_OF_INTEREST = EnumSet.of(Case.Events.ADDING_DATA_SOURCE,
            Case.Events.DATA_SOURCE_ADDED, Case.Events.ADDING_DATA_SOURCE_FAILED);
    private static final Set<IngestManager.IngestJobEvent> INGEST_JOB_EVENTS_OF_INTEREST = EnumSet.of(IngestManager.IngestJobEvent.DATA_SOURCE_ANALYSIS_STARTED, IngestManager.IngestJobEvent.DATA_SOURCE_ANALYSIS_COMPLETED);
    private static final int NUMBER_OF_PERIODIC_TASK_THREADS = 2;
    private static final String PERIODIC_TASK_THREAD_NAME = "collab-monitor-periodic-tasks-%d"; //NON-NLS
    private static final long HEARTBEAT_INTERVAL_MINUTES = 1;
    private static final long MAX_MISSED_HEARTBEATS = 5;
    private static final long STALE_TASKS_DETECT_INTERVAL_MINS = 2;
    private static final long EXECUTOR_TERMINATION_WAIT_SECS = 30;
    private static final Logger logger = Logger.getLogger(CollaborationMonitor.class.getName());
    private final String hostName;
    private final LocalTasksManager localTasksManager;
    private final RemoteTasksManager remoteTasksManager;
    private final AutopsyEventPublisher eventPublisher;
    private final ScheduledThreadPoolExecutor periodicTasksExecutor;

    CollaborationMonitor(String eventChannelPrefix) throws CollaborationMonitorException {
        hostName = NetworkUtils.getLocalHostName();

        eventPublisher = new AutopsyEventPublisher();
        try {
            eventPublisher.openRemoteEventChannel(String.format(EVENT_CHANNEL_NAME, eventChannelPrefix));
        } catch (AutopsyEventException ex) {
            throw new CollaborationMonitorException("Failed to initialize", ex);
        }

        remoteTasksManager = new RemoteTasksManager();
        eventPublisher.addSubscriber(COLLABORATION_MONITOR_EVENT, remoteTasksManager);

        localTasksManager = new LocalTasksManager();
        IngestManager.getInstance().addIngestJobEventListener(INGEST_JOB_EVENTS_OF_INTEREST, localTasksManager);
        Case.addEventTypeSubscriber(CASE_EVENTS_OF_INTEREST, localTasksManager);

        periodicTasksExecutor = new ScheduledThreadPoolExecutor(NUMBER_OF_PERIODIC_TASK_THREADS, new ThreadFactoryBuilder().setNameFormat(PERIODIC_TASK_THREAD_NAME).build());
        periodicTasksExecutor.scheduleWithFixedDelay(new HeartbeatTask(), HEARTBEAT_INTERVAL_MINUTES, HEARTBEAT_INTERVAL_MINUTES, TimeUnit.MINUTES);
        periodicTasksExecutor.scheduleWithFixedDelay(new StaleTaskDetectionTask(), STALE_TASKS_DETECT_INTERVAL_MINS, STALE_TASKS_DETECT_INTERVAL_MINS, TimeUnit.MINUTES);
    }

    void shutdown() {
        if (null != periodicTasksExecutor) {
            periodicTasksExecutor.shutdownNow();
            try {
                while (!periodicTasksExecutor.awaitTermination(EXECUTOR_TERMINATION_WAIT_SECS, TimeUnit.SECONDS)) {
                    logger.log(Level.WARNING, "Waited at least {0} seconds for periodic tasks executor to shut down, continuing to wait", EXECUTOR_TERMINATION_WAIT_SECS); //NON-NLS
                }
            } catch (InterruptedException ex) {
                logger.log(Level.SEVERE, "Unexpected interrupt while stopping periodic tasks executor", ex); //NON-NLS
            }
        }

        Case.removeEventTypeSubscriber(CASE_EVENTS_OF_INTEREST, localTasksManager);
        IngestManager.getInstance().removeIngestJobEventListener(localTasksManager);

        if (null != eventPublisher) {
            eventPublisher.removeSubscriber(COLLABORATION_MONITOR_EVENT, remoteTasksManager);
            eventPublisher.closeRemoteEventChannel();
        }

        remoteTasksManager.shutdown();
    }

    private final class LocalTasksManager implements PropertyChangeListener {

        private long nextTaskId;
        private final Map<Integer, Task> eventIdsToAddDataSourceTasks;
        private final Map<Long, Task> jobIdsTodataSourceAnalysisTasks;

        LocalTasksManager() {
            nextTaskId = 0;
            eventIdsToAddDataSourceTasks = new HashMap<>();
            jobIdsTodataSourceAnalysisTasks = new HashMap<>();
        }

        @Override
        public void propertyChange(PropertyChangeEvent event) {
            if (AutopsyEvent.SourceType.LOCAL == ((AutopsyEvent) event).getSourceType()) {
                String eventName = event.getPropertyName();
                if (eventName.equals(Case.Events.ADDING_DATA_SOURCE.toString())) {
                    addDataSourceAddTask((AddingDataSourceEvent) event);
                } else if (eventName.equals(Case.Events.ADDING_DATA_SOURCE_FAILED.toString())) {
                    removeDataSourceAddTask(((AddingDataSourceFailedEvent) event).getAddingDataSourceEventId());
                } else if (eventName.equals(Case.Events.DATA_SOURCE_ADDED.toString())) {
                    removeDataSourceAddTask(((DataSourceAddedEvent) event).getAddingDataSourceEventId());
                } else if (eventName.equals(IngestManager.IngestJobEvent.DATA_SOURCE_ANALYSIS_STARTED.toString())) {
                    addDataSourceAnalysisTask((DataSourceAnalysisStartedEvent) event);
                } else if (eventName.equals(IngestManager.IngestJobEvent.DATA_SOURCE_ANALYSIS_COMPLETED.toString())) {
                    removeDataSourceAnalysisTask((DataSourceAnalysisCompletedEvent) event);
                }
            }
        }

        synchronized void addDataSourceAddTask(AddingDataSourceEvent event) {
            String status = NbBundle.getMessage(CollaborationMonitor.class, "CollaborationMonitor.addingDataSourceStatus.msg", hostName);
            eventIdsToAddDataSourceTasks.put(event.getEventId().hashCode(), new Task(++nextTaskId, status));
            eventPublisher.publishRemotely(new CollaborationEvent(hostName, getCurrentTasks()));
        }

        synchronized void removeDataSourceAddTask(UUID eventId) {
            eventIdsToAddDataSourceTasks.remove(eventId.hashCode());
            eventPublisher.publishRemotely(new CollaborationEvent(hostName, getCurrentTasks()));
        }

        synchronized void addDataSourceAnalysisTask(DataSourceAnalysisStartedEvent event) {
            Content dataSource = event.getDataSource();
            if (dataSource != null) {
                String status = NbBundle.getMessage(CollaborationMonitor.class, "CollaborationMonitor.analyzingDataSourceStatus.msg", hostName, dataSource.getName());
                jobIdsTodataSourceAnalysisTasks.put(event.getIngestJobId(), new Task(++nextTaskId, status));
                eventPublisher.publishRemotely(new CollaborationEvent(hostName, getCurrentTasks()));
            }
        }

        synchronized void removeDataSourceAnalysisTask(DataSourceAnalysisCompletedEvent event) {
            jobIdsTodataSourceAnalysisTasks.remove(event.getIngestJobId());
            eventPublisher.publishRemotely(new CollaborationEvent(hostName, getCurrentTasks()));
        }

        synchronized Map<Long, Task> getCurrentTasks() {
            Map<Long, Task> currentTasks = new HashMap<>();
            eventIdsToAddDataSourceTasks.values().stream().forEach((task) -> {
                currentTasks.put(task.getId(), task);
            });
            jobIdsTodataSourceAnalysisTasks.values().stream().forEach((task) -> {
                currentTasks.put(task.getId(), task);
            });
            return currentTasks;
        }
    }

    private final class RemoteTasksManager implements PropertyChangeListener {

        private final Map<String, RemoteTasks> hostsToTasks;

        RemoteTasksManager() {
            hostsToTasks = new HashMap<>();
        }

        @Override
        public void propertyChange(PropertyChangeEvent event) {
            if (event.getPropertyName().equals(COLLABORATION_MONITOR_EVENT)) {
                updateTasks((CollaborationEvent) event);
            }
        }

        synchronized void shutdown() {
            finishAllTasks();
        }

        synchronized void updateTasks(CollaborationEvent event) {
            RemoteTasks tasksForHost = hostsToTasks.get(event.getHostName());
            if (null != tasksForHost) {
                tasksForHost.update(event);
            } else {
                hostsToTasks.put(event.getHostName(), new RemoteTasks(event));
            }
        }

        synchronized void finishStaleTasks() {
            for (Iterator<Map.Entry<String, RemoteTasks>> it = hostsToTasks.entrySet().iterator(); it.hasNext();) {
                Map.Entry<String, RemoteTasks> entry = it.next();
                RemoteTasks tasksForHost = entry.getValue();
                if (tasksForHost.isStale()) {
                    tasksForHost.finishAllTasks();
                    it.remove();
                }
            }
        }

        synchronized void finishAllTasks() {
            for (Iterator<Map.Entry<String, RemoteTasks>> it = hostsToTasks.entrySet().iterator(); it.hasNext();) {
                Map.Entry<String, RemoteTasks> entry = it.next();
                RemoteTasks tasksForHost = entry.getValue();
                tasksForHost.finishAllTasks();
                it.remove();
            }
        }

        private final class RemoteTasks {

            private final long MAX_MINUTES_WITHOUT_UPDATE = HEARTBEAT_INTERVAL_MINUTES * MAX_MISSED_HEARTBEATS;
            private Instant lastUpdateTime;
            private Map<Long, ProgressHandle> taskIdsToProgressBars;

            RemoteTasks(CollaborationEvent event) {
                lastUpdateTime = Instant.now();

                taskIdsToProgressBars = new HashMap<>();
                event.getCurrentTasks().values().stream().forEach((task) -> {
                    ProgressHandle progress = ProgressHandle.createHandle(event.getHostName());
                    progress.start();
                    progress.progress(task.getStatus());
                    taskIdsToProgressBars.put(task.getId(), progress);
                });
            }

            void update(CollaborationEvent event) {
                lastUpdateTime = Instant.now();

                Map<Long, Task> remoteTasks = event.getCurrentTasks();
                remoteTasks.values().stream().forEach((task) -> {
                    ProgressHandle progress = taskIdsToProgressBars.get(task.getId());
                    if (null != progress) {
                        progress.progress(task.getStatus());
                    } else {
                        progress = ProgressHandle.createHandle(event.getHostName());
                        progress.start();
                        progress.progress(task.getStatus());
                        taskIdsToProgressBars.put(task.getId(), progress);
                    }
                });

                for (Iterator<Map.Entry<Long, ProgressHandle>> iterator = taskIdsToProgressBars.entrySet().iterator(); iterator.hasNext();) {
                    Map.Entry<Long, ProgressHandle> entry = iterator.next();
                    if (!remoteTasks.containsKey(entry.getKey())) {
                        ProgressHandle progress = entry.getValue();
                        progress.finish();
                        iterator.remove();
                    }
                }
            }

            void finishAllTasks() {
                taskIdsToProgressBars.values().stream().forEach((progress) -> {
                    progress.finish();
                });
                taskIdsToProgressBars.clear();
            }

            boolean isStale() {
                return Duration.between(lastUpdateTime, Instant.now()).toMinutes() >= MAX_MINUTES_WITHOUT_UPDATE;
            }
        }

    }

    private final class HeartbeatTask implements Runnable {

        @Override
        public void run() {
            try {
                eventPublisher.publishRemotely(new CollaborationEvent(hostName, localTasksManager.getCurrentTasks()));
            } catch (Exception ex) {
                logger.log(Level.SEVERE, "Unexpected exception in HeartbeatTask", ex); //NON-NLS
            }
        }
    }

    private final class StaleTaskDetectionTask implements Runnable {

        @Override
        public void run() {
            try {
                remoteTasksManager.finishStaleTasks();
            } catch (Exception ex) {
                logger.log(Level.SEVERE, "Unexpected exception in StaleTaskDetectionTask", ex); //NON-NLS
            }
        }
    }

    private final static class CollaborationEvent extends AutopsyEvent implements Serializable {

        private static final long serialVersionUID = 1L;
        private final String hostName;
        private final Map<Long, Task> currentTasks;

        CollaborationEvent(String hostName, Map<Long, Task> currentTasks) {
            super(COLLABORATION_MONITOR_EVENT, null, null);
            this.hostName = hostName;
            this.currentTasks = currentTasks;
        }

        String getHostName() {
            return hostName;
        }

        Map<Long, Task> getCurrentTasks() {
            return Collections.unmodifiableMap(currentTasks);
        }

    }

    private final static class Task implements Serializable {

        private static final long serialVersionUID = 1L;
        private final long id;
        private final String status;

        Task(long id, String status) {
            this.id = id;
            this.status = status;
        }

        long getId() {
            return id;
        }

        String getStatus() {
            return status;
        }
    }

    final static class CollaborationMonitorException extends Exception {

        CollaborationMonitorException(String message) {
            super(message);
        }

        CollaborationMonitorException(String message, Throwable throwable) {
            super(message, throwable);
        }
    }

}