Distributed work queue primitive
Change-Id: Ia8e531e6611ec502399edec376ccc00522e47994
diff --git a/core/store/primitives/src/main/java/org/onosproject/store/primitives/resources/impl/AtomixWorkQueue.java b/core/store/primitives/src/main/java/org/onosproject/store/primitives/resources/impl/AtomixWorkQueue.java
new file mode 100644
index 0000000..7b4ad47
--- /dev/null
+++ b/core/store/primitives/src/main/java/org/onosproject/store/primitives/resources/impl/AtomixWorkQueue.java
@@ -0,0 +1,201 @@
+/*
+ * Copyright 2016-present Open Networking Laboratory
+ *
+ * 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.onosproject.store.primitives.resources.impl;
+
+import static org.slf4j.LoggerFactory.getLogger;
+
+import java.util.Collection;
+import java.util.List;
+import java.util.Properties;
+import java.util.Timer;
+import java.util.concurrent.CompletableFuture;
+import java.util.concurrent.Executor;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.atomic.AtomicBoolean;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Consumer;
+
+import org.onlab.util.AbstractAccumulator;
+import org.onlab.util.Accumulator;
+import org.onosproject.store.primitives.resources.impl.AtomixWorkQueueCommands.Add;
+import org.onosproject.store.primitives.resources.impl.AtomixWorkQueueCommands.Complete;
+import org.onosproject.store.primitives.resources.impl.AtomixWorkQueueCommands.Register;
+import org.onosproject.store.primitives.resources.impl.AtomixWorkQueueCommands.Stats;
+import org.onosproject.store.primitives.resources.impl.AtomixWorkQueueCommands.Take;
+import org.onosproject.store.primitives.resources.impl.AtomixWorkQueueCommands.Unregister;
+import org.onosproject.store.service.WorkQueue;
+import org.onosproject.store.service.Task;
+import org.onosproject.store.service.WorkQueueStats;
+import org.slf4j.Logger;
+
+import com.google.common.collect.ImmutableList;
+
+import io.atomix.copycat.client.CopycatClient;
+import io.atomix.resource.AbstractResource;
+import io.atomix.resource.ResourceTypeInfo;
+
+/**
+ * Distributed resource providing the {@link WorkQueue} primitive.
+ */
+@ResourceTypeInfo(id = -154, factory = AtomixWorkQueueFactory.class)
+public class AtomixWorkQueue extends AbstractResource<AtomixWorkQueue>
+ implements WorkQueue<byte[]> {
+
+ private final Logger log = getLogger(getClass());
+ public static final String TASK_AVAILABLE = "task-available";
+ private final ExecutorService executor = Executors.newSingleThreadExecutor();
+ private final AtomicReference<TaskProcessor> taskProcessor = new AtomicReference<>();
+ private final Timer timer = new Timer("atomix-work-queue-completer");
+ private final AtomicBoolean isRegistered = new AtomicBoolean(false);
+
+ protected AtomixWorkQueue(CopycatClient client, Properties options) {
+ super(client, options);
+ }
+
+ @Override
+ public CompletableFuture<AtomixWorkQueue> open() {
+ return super.open().thenApply(result -> {
+ client.onStateChange(state -> {
+ if (state == CopycatClient.State.CONNECTED && isRegistered.get()) {
+ client.submit(new Register());
+ }
+ });
+ client.onEvent(TASK_AVAILABLE, this::resumeWork);
+ return result;
+ });
+ }
+
+ @Override
+ public CompletableFuture<Void> addMultiple(Collection<byte[]> items) {
+ if (items.isEmpty()) {
+ return CompletableFuture.completedFuture(null);
+ }
+ return client.submit(new Add(items));
+ }
+
+ @Override
+ public CompletableFuture<Collection<Task<byte[]>>> take(int maxTasks) {
+ if (maxTasks <= 0) {
+ return CompletableFuture.completedFuture(ImmutableList.of());
+ }
+ return client.submit(new Take(maxTasks));
+ }
+
+ @Override
+ public CompletableFuture<Void> complete(Collection<String> taskIds) {
+ if (taskIds.isEmpty()) {
+ return CompletableFuture.completedFuture(null);
+ }
+ return client.submit(new Complete(taskIds));
+ }
+
+ @Override
+ public CompletableFuture<Void> registerTaskProcessor(Consumer<byte[]> callback,
+ int parallelism,
+ Executor executor) {
+ Accumulator<String> completedTaskAccumulator =
+ new CompletedTaskAccumulator(timer, 50, 50); // TODO: make configurable
+ taskProcessor.set(new TaskProcessor(callback,
+ parallelism,
+ executor,
+ completedTaskAccumulator));
+ return register().thenCompose(v -> take(parallelism))
+ .thenAccept(taskProcessor.get()::accept);
+ }
+
+ @Override
+ public CompletableFuture<Void> stopProcessing() {
+ return unregister();
+ }
+
+ @Override
+ public CompletableFuture<WorkQueueStats> stats() {
+ return client.submit(new Stats());
+ }
+
+ private void resumeWork() {
+ TaskProcessor activeProcessor = taskProcessor.get();
+ if (activeProcessor == null) {
+ return;
+ }
+ this.take(activeProcessor.headRoom())
+ .whenCompleteAsync((tasks, e) -> activeProcessor.accept(tasks), executor);
+ }
+
+ private CompletableFuture<Void> register() {
+ return client.submit(new Register()).thenRun(() -> isRegistered.set(true));
+ }
+
+ private CompletableFuture<Void> unregister() {
+ return client.submit(new Unregister()).thenRun(() -> isRegistered.set(false));
+ }
+
+ // TaskId accumulator for paced triggering of task completion calls.
+ private class CompletedTaskAccumulator extends AbstractAccumulator<String> {
+ CompletedTaskAccumulator(Timer timer, int maxTasksToBatch, int maxBatchMillis) {
+ super(timer, maxTasksToBatch, maxBatchMillis, Integer.MAX_VALUE);
+ }
+
+ @Override
+ public void processItems(List<String> items) {
+ complete(items);
+ }
+ }
+
+ private class TaskProcessor implements Consumer<Collection<Task<byte[]>>> {
+
+ private final AtomicInteger headRoom;
+ private final Consumer<byte[]> backingConsumer;
+ private final Executor executor;
+ private final Accumulator<String> taskCompleter;
+
+ public TaskProcessor(Consumer<byte[]> backingConsumer,
+ int parallelism,
+ Executor executor,
+ Accumulator<String> taskCompleter) {
+ this.backingConsumer = backingConsumer;
+ this.headRoom = new AtomicInteger(parallelism);
+ this.executor = executor;
+ this.taskCompleter = taskCompleter;
+ }
+
+ public int headRoom() {
+ return headRoom.get();
+ }
+
+ @Override
+ public void accept(Collection<Task<byte[]>> tasks) {
+ if (tasks == null) {
+ return;
+ }
+ headRoom.addAndGet(-1 * tasks.size());
+ tasks.forEach(task ->
+ executor.execute(() -> {
+ try {
+ backingConsumer.accept(task.payload());
+ taskCompleter.add(task.taskId());
+ } catch (Exception e) {
+ log.debug("Task execution failed", e);
+ } finally {
+ headRoom.incrementAndGet();
+ resumeWork();
+ }
+ }));
+ }
+ }
+}
diff --git a/core/store/primitives/src/main/java/org/onosproject/store/primitives/resources/impl/AtomixWorkQueueCommands.java b/core/store/primitives/src/main/java/org/onosproject/store/primitives/resources/impl/AtomixWorkQueueCommands.java
new file mode 100644
index 0000000..3724529
--- /dev/null
+++ b/core/store/primitives/src/main/java/org/onosproject/store/primitives/resources/impl/AtomixWorkQueueCommands.java
@@ -0,0 +1,224 @@
+/*
+ * Copyright 2016-present Open Networking Laboratory
+ *
+ * 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.onosproject.store.primitives.resources.impl;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import org.onosproject.store.service.Task;
+import org.onosproject.store.service.WorkQueueStats;
+
+import com.google.common.base.MoreObjects;
+
+import io.atomix.catalyst.buffer.BufferInput;
+import io.atomix.catalyst.buffer.BufferOutput;
+import io.atomix.catalyst.serializer.CatalystSerializable;
+import io.atomix.catalyst.serializer.SerializableTypeResolver;
+import io.atomix.catalyst.serializer.Serializer;
+import io.atomix.catalyst.serializer.SerializerRegistry;
+import io.atomix.copycat.Command;
+
+/**
+ * {@link AtomixWorkQueue} resource state machine operations.
+ */
+public final class AtomixWorkQueueCommands {
+
+ private AtomixWorkQueueCommands() {
+ }
+
+ /**
+ * Command to add a collection of tasks to the queue.
+ */
+ @SuppressWarnings("serial")
+ public static class Add implements Command<Void>, CatalystSerializable {
+
+ private Collection<byte[]> items;
+
+ private Add() {
+ }
+
+ public Add(Collection<byte[]> items) {
+ this.items = items;
+ }
+
+ @Override
+ public void writeObject(BufferOutput<?> buffer, Serializer serializer) {
+ buffer.writeInt(items.size());
+ items.forEach(task -> serializer.writeObject(task, buffer));
+ }
+
+ @Override
+ public void readObject(BufferInput<?> buffer, Serializer serializer) {
+ items = IntStream.range(0, buffer.readInt())
+ .mapToObj(i -> serializer.<byte[]>readObject(buffer))
+ .collect(Collectors.toCollection(ArrayList::new));
+ }
+
+ public Collection<byte[]> items() {
+ return items;
+ }
+
+ @Override
+ public String toString() {
+ return MoreObjects.toStringHelper(getClass())
+ .add("items", items)
+ .toString();
+ }
+ }
+
+ /**
+ * Command to take a task from the queue.
+ */
+ @SuppressWarnings("serial")
+ public static class Take implements Command<Collection<Task<byte[]>>>, CatalystSerializable {
+
+ private int maxTasks;
+
+ private Take() {
+ }
+
+ public Take(int maxTasks) {
+ this.maxTasks = maxTasks;
+ }
+
+ @Override
+ public void writeObject(BufferOutput<?> buffer, Serializer serializer) {
+ buffer.writeInt(maxTasks);
+ }
+
+ @Override
+ public void readObject(BufferInput<?> buffer, Serializer serializer) {
+ maxTasks = buffer.readInt();
+ }
+
+ public int maxTasks() {
+ return maxTasks;
+ }
+
+ @Override
+ public String toString() {
+ return MoreObjects.toStringHelper(getClass())
+ .add("maxTasks", maxTasks)
+ .toString();
+ }
+ }
+
+ @SuppressWarnings("serial")
+ public static class Stats implements Command<WorkQueueStats>, CatalystSerializable {
+
+ @Override
+ public void writeObject(BufferOutput<?> buffer, Serializer serializer) {
+ }
+
+ @Override
+ public void readObject(BufferInput<?> buffer, Serializer serializer) {
+ }
+
+ @Override
+ public String toString() {
+ return MoreObjects.toStringHelper(getClass())
+ .toString();
+ }
+ }
+
+
+
+ @SuppressWarnings("serial")
+ public static class Register implements Command<Void>, CatalystSerializable {
+
+ @Override
+ public void writeObject(BufferOutput<?> buffer, Serializer serializer) {
+ }
+
+ @Override
+ public void readObject(BufferInput<?> buffer, Serializer serializer) {
+ }
+
+ @Override
+ public String toString() {
+ return MoreObjects.toStringHelper(getClass())
+ .toString();
+ }
+ }
+
+ @SuppressWarnings("serial")
+ public static class Unregister implements Command<Void>, CatalystSerializable {
+
+ @Override
+ public void writeObject(BufferOutput<?> buffer, Serializer serializer) {
+ }
+
+ @Override
+ public void readObject(BufferInput<?> buffer, Serializer serializer) {
+ }
+
+ @Override
+ public String toString() {
+ return MoreObjects.toStringHelper(getClass())
+ .toString();
+ }
+ }
+
+ @SuppressWarnings("serial")
+ public static class Complete implements Command<Void>, CatalystSerializable {
+ private Collection<String> taskIds;
+
+ private Complete() {
+ }
+
+ public Complete(Collection<String> taskIds) {
+ this.taskIds = taskIds;
+ }
+
+ @Override
+ public void writeObject(BufferOutput<?> buffer, Serializer serializer) {
+ serializer.writeObject(taskIds, buffer);
+ }
+
+ @Override
+ public void readObject(BufferInput<?> buffer, Serializer serializer) {
+ taskIds = serializer.readObject(buffer);
+ }
+
+ public Collection<String> taskIds() {
+ return taskIds;
+ }
+
+ @Override
+ public String toString() {
+ return MoreObjects.toStringHelper(getClass())
+ .add("taskIds", taskIds)
+ .toString();
+ }
+ }
+
+ /**
+ * Work queue command type resolver.
+ */
+ public static class TypeResolver implements SerializableTypeResolver {
+ @Override
+ public void resolve(SerializerRegistry registry) {
+ registry.register(Register.class, -960);
+ registry.register(Unregister.class, -961);
+ registry.register(Take.class, -962);
+ registry.register(Add.class, -963);
+ registry.register(Complete.class, -964);
+ registry.register(Stats.class, -965);
+ }
+ }
+}
diff --git a/core/store/primitives/src/main/java/org/onosproject/store/primitives/resources/impl/AtomixWorkQueueFactory.java b/core/store/primitives/src/main/java/org/onosproject/store/primitives/resources/impl/AtomixWorkQueueFactory.java
new file mode 100644
index 0000000..0c61b2e
--- /dev/null
+++ b/core/store/primitives/src/main/java/org/onosproject/store/primitives/resources/impl/AtomixWorkQueueFactory.java
@@ -0,0 +1,44 @@
+/*
+ * Copyright 2016-present Open Networking Laboratory
+ *
+ * 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.onosproject.store.primitives.resources.impl;
+
+import io.atomix.catalyst.serializer.SerializableTypeResolver;
+import io.atomix.copycat.client.CopycatClient;
+import io.atomix.resource.ResourceFactory;
+import io.atomix.resource.ResourceStateMachine;
+
+import java.util.Properties;
+
+/**
+ * {@link AtomixWorkQueue} resource factory.
+ */
+public class AtomixWorkQueueFactory implements ResourceFactory<AtomixWorkQueue> {
+
+ @Override
+ public SerializableTypeResolver createSerializableTypeResolver() {
+ return new AtomixWorkQueueCommands.TypeResolver();
+ }
+
+ @Override
+ public ResourceStateMachine createStateMachine(Properties config) {
+ return new AtomixWorkQueueState(config);
+ }
+
+ @Override
+ public AtomixWorkQueue createInstance(CopycatClient client, Properties properties) {
+ return new AtomixWorkQueue(client, properties);
+ }
+}
\ No newline at end of file
diff --git a/core/store/primitives/src/main/java/org/onosproject/store/primitives/resources/impl/AtomixWorkQueueState.java b/core/store/primitives/src/main/java/org/onosproject/store/primitives/resources/impl/AtomixWorkQueueState.java
new file mode 100644
index 0000000..d287e19
--- /dev/null
+++ b/core/store/primitives/src/main/java/org/onosproject/store/primitives/resources/impl/AtomixWorkQueueState.java
@@ -0,0 +1,289 @@
+/*
+ * Copyright 2016-present Open Networking Laboratory
+ *
+ * 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.onosproject.store.primitives.resources.impl;
+
+import static org.slf4j.LoggerFactory.getLogger;
+
+import java.util.ArrayList;
+import java.util.Collection;
+import java.util.Iterator;
+import java.util.Map;
+import java.util.Properties;
+import java.util.Queue;
+import java.util.concurrent.atomic.AtomicInteger;
+import java.util.concurrent.atomic.AtomicLong;
+import java.util.function.Consumer;
+import java.util.stream.Collectors;
+import java.util.stream.IntStream;
+
+import org.onlab.util.CountDownCompleter;
+import org.onosproject.store.primitives.resources.impl.AtomixWorkQueueCommands.Add;
+import org.onosproject.store.primitives.resources.impl.AtomixWorkQueueCommands.Complete;
+import org.onosproject.store.primitives.resources.impl.AtomixWorkQueueCommands.Register;
+import org.onosproject.store.primitives.resources.impl.AtomixWorkQueueCommands.Stats;
+import org.onosproject.store.primitives.resources.impl.AtomixWorkQueueCommands.Take;
+import org.onosproject.store.primitives.resources.impl.AtomixWorkQueueCommands.Unregister;
+import org.onosproject.store.service.Task;
+import org.onosproject.store.service.WorkQueueStats;
+import org.slf4j.Logger;
+
+import com.google.common.base.MoreObjects;
+import com.google.common.base.Throwables;
+import com.google.common.collect.ImmutableList;
+import com.google.common.collect.Maps;
+import com.google.common.collect.Queues;
+import com.google.common.util.concurrent.AtomicLongMap;
+
+import io.atomix.copycat.server.Commit;
+import io.atomix.copycat.server.Snapshottable;
+import io.atomix.copycat.server.StateMachineExecutor;
+import io.atomix.copycat.server.session.ServerSession;
+import io.atomix.copycat.server.session.SessionListener;
+import io.atomix.copycat.server.storage.snapshot.SnapshotReader;
+import io.atomix.copycat.server.storage.snapshot.SnapshotWriter;
+import io.atomix.resource.ResourceStateMachine;
+
+/**
+ * State machine for {@link AtomixWorkQueue} resource.
+ */
+public class AtomixWorkQueueState extends ResourceStateMachine implements SessionListener, Snapshottable {
+
+ private final Logger log = getLogger(getClass());
+
+ private final AtomicLong totalCompleted = new AtomicLong(0);
+
+ private final Queue<TaskHolder> unassignedTasks = Queues.newArrayDeque();
+ private final Map<String, TaskAssignment> assignments = Maps.newHashMap();
+ private final Map<Long, Commit<? extends Register>> registeredWorkers = Maps.newHashMap();
+ private final AtomicLongMap<Long> activeTasksPerSession = AtomicLongMap.create();
+
+ protected AtomixWorkQueueState(Properties config) {
+ super(config);
+ }
+
+ @Override
+ protected void configure(StateMachineExecutor executor) {
+ executor.register(Stats.class, this::stats);
+ executor.register(Register.class, (Consumer<Commit<Register>>) this::register);
+ executor.register(Unregister.class, (Consumer<Commit<Unregister>>) this::unregister);
+ executor.register(Add.class, (Consumer<Commit<Add>>) this::add);
+ executor.register(Take.class, this::take);
+ executor.register(Complete.class, (Consumer<Commit<Complete>>) this::complete);
+ }
+
+ protected WorkQueueStats stats(Commit<? extends Stats> commit) {
+ try {
+ return WorkQueueStats.builder()
+ .withTotalCompleted(totalCompleted.get())
+ .withTotalPending(unassignedTasks.size())
+ .withTotalInProgress(assignments.size())
+ .build();
+ } finally {
+ commit.close();
+ }
+ }
+
+ protected void register(Commit<? extends Register> commit) {
+ long sessionId = commit.session().id();
+ if (registeredWorkers.putIfAbsent(sessionId, commit) != null) {
+ commit.close();
+ }
+ }
+
+ protected void unregister(Commit<? extends Unregister> commit) {
+ try {
+ Commit<? extends Register> registerCommit = registeredWorkers.remove(commit.session().id());
+ if (registerCommit != null) {
+ registerCommit.close();
+ }
+ } finally {
+ commit.close();
+ }
+ }
+
+ protected void add(Commit<? extends Add> commit) {
+ Collection<byte[]> items = commit.operation().items();
+
+ // Create a CountDownCompleter that will close the commit when all tasks
+ // submitted as part of it are completed.
+ CountDownCompleter<Commit<? extends Add>> referenceTracker =
+ new CountDownCompleter<>(commit, items.size(), Commit::close);
+
+ AtomicInteger itemIndex = new AtomicInteger(0);
+ items.forEach(item -> {
+ String taskId = String.format("%d:%d:%d", commit.session().id(),
+ commit.index(),
+ itemIndex.getAndIncrement());
+ unassignedTasks.add(new TaskHolder(new Task<>(taskId, item), referenceTracker));
+ });
+
+ // Send an event to all sessions that have expressed interest in task processing
+ // and are not actively processing a task.
+ registeredWorkers.values()
+ .stream()
+ .map(Commit::session)
+ .forEach(session -> session.publish(AtomixWorkQueue.TASK_AVAILABLE));
+ // FIXME: This generates a lot of event traffic.
+ }
+
+ protected Collection<Task<byte[]>> take(Commit<? extends Take> commit) {
+ try {
+ if (unassignedTasks.isEmpty()) {
+ return ImmutableList.of();
+ }
+ long sessionId = commit.session().id();
+ int maxTasks = commit.operation().maxTasks();
+ return IntStream.range(0, Math.min(maxTasks, unassignedTasks.size()))
+ .mapToObj(i -> {
+ TaskHolder holder = unassignedTasks.poll();
+ String taskId = holder.task().taskId();
+ TaskAssignment assignment = new TaskAssignment(sessionId, holder);
+
+ // bookkeeping
+ assignments.put(taskId, assignment);
+ activeTasksPerSession.incrementAndGet(sessionId);
+
+ return holder.task();
+ })
+ .collect(Collectors.toCollection(ArrayList::new));
+ } catch (Exception e) {
+ log.warn("State machine update failed", e);
+ throw Throwables.propagate(e);
+ } finally {
+ commit.close();
+ }
+ }
+
+ protected void complete(Commit<? extends Complete> commit) {
+ long sessionId = commit.session().id();
+ try {
+ commit.operation().taskIds().forEach(taskId -> {
+ TaskAssignment assignment = assignments.get(taskId);
+ if (assignment != null) {
+ assignments.remove(taskId).markComplete();
+ // bookkeeping
+ totalCompleted.incrementAndGet();
+ activeTasksPerSession.decrementAndGet(sessionId);
+ }
+ });
+ } catch (Exception e) {
+ log.warn("State machine update failed", e);
+ throw Throwables.propagate(e);
+ } finally {
+ commit.close();
+ }
+ }
+
+ @Override
+ public void register(ServerSession session) {
+ }
+
+ @Override
+ public void unregister(ServerSession session) {
+ evictWorker(session.id());
+ }
+
+ @Override
+ public void expire(ServerSession session) {
+ evictWorker(session.id());
+ }
+
+ @Override
+ public void close(ServerSession session) {
+ evictWorker(session.id());
+ }
+
+ @Override
+ public void snapshot(SnapshotWriter writer) {
+ writer.writeLong(totalCompleted.get());
+ }
+
+ @Override
+ public void install(SnapshotReader reader) {
+ totalCompleted.set(reader.readLong());
+ }
+
+ private void evictWorker(long sessionId) {
+ Commit<? extends Register> commit = registeredWorkers.remove(sessionId);
+ if (commit != null) {
+ commit.close();
+ }
+
+ // TODO: Maintain an index of tasks by session for efficient access.
+ Iterator<Map.Entry<String, TaskAssignment>> iter = assignments.entrySet().iterator();
+ while (iter.hasNext()) {
+ Map.Entry<String, TaskAssignment> entry = iter.next();
+ TaskAssignment assignment = entry.getValue();
+ if (assignment.sessionId() == sessionId) {
+ unassignedTasks.add(assignment.taskHolder());
+ iter.remove();
+ }
+ }
+
+ // Bookkeeping
+ activeTasksPerSession.remove(sessionId);
+ activeTasksPerSession.removeAllZeros();
+ }
+
+ private class TaskHolder {
+
+ private final Task<byte[]> task;
+ private final CountDownCompleter<Commit<? extends Add>> referenceTracker;
+
+ public TaskHolder(Task<byte[]> delegate, CountDownCompleter<Commit<? extends Add>> referenceTracker) {
+ this.task = delegate;
+ this.referenceTracker = referenceTracker;
+ }
+
+ public Task<byte[]> task() {
+ return task;
+ }
+
+ public void complete() {
+ referenceTracker.countDown();
+ }
+ }
+
+ private class TaskAssignment {
+ private final long sessionId;
+ private final TaskHolder taskHolder;
+
+ public TaskAssignment(long sessionId, TaskHolder taskHolder) {
+ this.sessionId = sessionId;
+ this.taskHolder = taskHolder;
+ }
+
+ public long sessionId() {
+ return sessionId;
+ }
+
+ public TaskHolder taskHolder() {
+ return taskHolder;
+ }
+
+ public void markComplete() {
+ taskHolder.complete();
+ }
+
+ @Override
+ public String toString() {
+ return MoreObjects.toStringHelper(getClass())
+ .add("sessionId", sessionId)
+ .add("taskHolder", taskHolder)
+ .toString();
+ }
+ }
+}