From 779582d4410b4bf612aa03bb17f0f951da1b5821 Mon Sep 17 00:00:00 2001
From: md_5 <md_5@live.com.au>
Date: Fri, 3 May 2013 14:33:04 +1000
Subject: [PATCH] Use multimap in scheduler

---
 .../net/md_5/bungee/scheduler/BungeeScheduler.java | 14 ++++++++------
 1 file changed, 8 insertions(+), 6 deletions(-)

diff --git a/proxy/src/main/java/net/md_5/bungee/scheduler/BungeeScheduler.java b/proxy/src/main/java/net/md_5/bungee/scheduler/BungeeScheduler.java
index fc7e73dda..2bf8b3901 100644
--- a/proxy/src/main/java/net/md_5/bungee/scheduler/BungeeScheduler.java
+++ b/proxy/src/main/java/net/md_5/bungee/scheduler/BungeeScheduler.java
@@ -1,6 +1,10 @@
 package net.md_5.bungee.scheduler;
 
 import com.google.common.base.Preconditions;
+import com.google.common.collect.ArrayListMultimap;
+import com.google.common.collect.HashMultimap;
+import com.google.common.collect.Multimap;
+import com.google.common.collect.Multimaps;
 import gnu.trove.TCollections;
 import gnu.trove.map.TIntObjectMap;
 import gnu.trove.map.hash.TIntObjectHashMap;
@@ -18,11 +22,13 @@ public class BungeeScheduler implements TaskScheduler
 
     private final AtomicInteger taskCounter = new AtomicInteger();
     private final TIntObjectMap<BungeeTask> tasks = TCollections.synchronizedMap( new TIntObjectHashMap<BungeeTask>() );
+    private final Multimap<Plugin, BungeeTask> tasksByPlugin = Multimaps.synchronizedMultimap( HashMultimap.<Plugin, BungeeTask>create() );
 
     @Override
     public void cancel(int id)
     {
         BungeeTask task = tasks.remove( id );
+        tasksByPlugin.values().remove( task );
         task.getFuture().cancel( false );
     }
 
@@ -36,13 +42,9 @@ public class BungeeScheduler implements TaskScheduler
     public int cancel(Plugin plugin)
     {
         Set<ScheduledTask> toRemove = new HashSet<>();
-        for ( ScheduledTask task : tasks.valueCollection() )
+        for ( ScheduledTask task : tasksByPlugin.get( plugin ) )
         {
-            // TODO: proper checking?
-            if ( task.getOwner() == plugin )
-            {
-                toRemove.add( task );
-            }
+            toRemove.add( task );
         }
         for ( ScheduledTask task : toRemove )
         {