package net.md_5.bungee.scheduler;

import com.google.common.base.Preconditions;
import gnu.trove.TCollections;
import gnu.trove.map.TIntObjectMap;
import gnu.trove.map.hash.TIntObjectHashMap;
import java.util.HashSet;
import java.util.Set;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;
import net.md_5.bungee.BungeeCord;
import net.md_5.bungee.api.plugin.Plugin;
import net.md_5.bungee.api.scheduler.ScheduledTask;
import net.md_5.bungee.api.scheduler.TaskScheduler;

public class BungeeScheduler implements TaskScheduler
{

    private final AtomicInteger taskCounter = new AtomicInteger();
    private final TIntObjectMap<BungeeTask> tasks = TCollections.synchronizedMap( new TIntObjectHashMap<BungeeTask>() );

    @Override
    public void cancel(int id)
    {
        BungeeTask task = tasks.remove( id );
        task.getFuture().cancel( false );
    }

    @Override
    public void cancel(ScheduledTask task)
    {
        cancel( task.getId() );
    }

    @Override
    public int cancel(Plugin plugin)
    {
        Set<ScheduledTask> toRemove = new HashSet<>();
        for ( ScheduledTask task : tasks.valueCollection() )
        {
            // TODO: proper checking?
            if ( task.getOwner() == plugin )
            {
                toRemove.add( task );
            }
        }
        for ( ScheduledTask task : toRemove )
        {
            cancel( task );
        }
        return toRemove.size();
    }

    @Override
    public ScheduledTask runAsync(Plugin owner, Runnable task)
    {
        return schedule( owner, task, 0, TimeUnit.MILLISECONDS );
    }

    @Override
    public ScheduledTask schedule(Plugin owner, Runnable task, long delay, TimeUnit unit)
    {
        return prepare( owner, task ).setFuture( BungeeCord.getInstance().executors.schedule( task, delay, unit ) );
    }

    @Override
    public ScheduledTask schedule(Plugin owner, Runnable task, long delay, long period, TimeUnit unit)
    {
        return prepare( owner, task ).setFuture( BungeeCord.getInstance().executors.scheduleWithFixedDelay( task, delay, period, unit ) );
    }

    private BungeeTask prepare(Plugin owner, Runnable task)
    {
        Preconditions.checkNotNull( owner, "owner" );
        Preconditions.checkNotNull( task, "task" );
        BungeeTask prepared = new BungeeTask( taskCounter.getAndIncrement(), owner, task );
        tasks.put( prepared.getId(), prepared );
        return prepared;
    }
}

@@ -70,24 +70,210 @@
 conditions;%0A
+import com.google.common.collect.ArrayListMultimap;%0Aimport com.google.common.collect.HashMultimap;%0Aimport com.google.common.collect.Multimap;%0Aimport com.google.common.collect.Multimaps;%0A
 import gnu.t
@@ -934,16 +934,156 @@
 sk%3E() );
+%0A    private final Multimap%3CPlugin, BungeeTask%3E tasksByPlugin = Multimaps.synchronizedMultimap( HashMultimap.%3CPlugin, BungeeTask%3Ecreate() );
 %0A%0A    @O
@@ -1169,24 +1169,71 @@
 move( id );%0A
+        tasksByPlugin.values().remove( task );%0A
         task
@@ -1527,112 +1527,21 @@
 asks
-.valueCollection() )%0A        %7B%0A            // TODO: proper checking?%0A            if ( task.getOwner() ==
+ByPlugin.get(
  plu
@@ -1545,16 +1545,18 @@
 plugin )
+ )
 %0A       
@@ -1552,33 +1552,26 @@
 ) )%0A        
-    %7B%0A   
+%7B%0A
             
@@ -1570,17 +1570,16 @@
         
- 
 toRemove
@@ -1596,22 +1596,8 @@
  );%0A
-            %7D%0A
     

Use multimap in scheduler