azkaban-aplcache

Preserve plugin order by using TreeSets in PluginRegistry

1/17/2014 2:20:10 AM

Details

diff --git a/src/java/azkaban/webapp/plugin/PluginRegistry.java b/src/java/azkaban/webapp/plugin/PluginRegistry.java
index 63e0726..0d98f32 100644
--- a/src/java/azkaban/webapp/plugin/PluginRegistry.java
+++ b/src/java/azkaban/webapp/plugin/PluginRegistry.java
@@ -20,29 +20,31 @@ import java.util.List;
 import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.Map;
+import java.util.TreeSet;
+import java.util.Set;
 
 public class PluginRegistry {
 
   private static PluginRegistry registry;
 
-  public Map<String, ViewerPlugin> viewerPlugins;
+  public TreeSet<ViewerPlugin> viewerPlugins;
 
-	public Map<String, List<ViewerPlugin>> jobTypeViewerPlugins;
+	public Map<String, TreeSet<ViewerPlugin>> jobTypeViewerPlugins;
 
   private PluginRegistry() {
-		viewerPlugins = new HashMap<String, ViewerPlugin>();
-		jobTypeViewerPlugins = new HashMap<String, List<ViewerPlugin>>();
+		viewerPlugins = new TreeSet<ViewerPlugin>(ViewerPlugin.COMPARATOR);
+		jobTypeViewerPlugins = new HashMap<String, TreeSet<ViewerPlugin>>();
   }
 
   public void register(ViewerPlugin plugin) {
-		viewerPlugins.put(plugin.getPluginName(), plugin);
+		viewerPlugins.add(plugin);
 		String jobType = plugin.getJobType();
 		if (jobType == null) {
 			return;
 		}
-		List<ViewerPlugin> plugins = null;
+		TreeSet<ViewerPlugin> plugins = null;
 		if (!jobTypeViewerPlugins.containsKey(jobType)) {
-			plugins = new ArrayList<ViewerPlugin>();
+			plugins = new TreeSet<ViewerPlugin>(ViewerPlugin.COMPARATOR);
 			plugins.add(plugin);
 			jobTypeViewerPlugins.put(jobType, plugins);
 		}
@@ -53,11 +55,11 @@ public class PluginRegistry {
   }
 
 	public List<ViewerPlugin> getViewerPlugins() {
-		return new ArrayList<ViewerPlugin>(viewerPlugins.values());
+		return new ArrayList<ViewerPlugin>(viewerPlugins);
 	}
 
 	public List<ViewerPlugin> getViewerPluginsForJobType(String jobType) {
-		return jobTypeViewerPlugins.get(jobType);
+		return new ArrayList<ViewerPlugin>(jobTypeViewerPlugins.get(jobType));
 	}
 
   public static PluginRegistry getRegistry() {
diff --git a/src/java/azkaban/webapp/plugin/ViewerPlugin.java b/src/java/azkaban/webapp/plugin/ViewerPlugin.java
index 4fd5fce..1ea864b 100644
--- a/src/java/azkaban/webapp/plugin/ViewerPlugin.java
+++ b/src/java/azkaban/webapp/plugin/ViewerPlugin.java
@@ -16,13 +16,23 @@
 
 package azkaban.webapp.plugin;
 
+import java.util.Comparator;
+
 public class ViewerPlugin {
 	private final String pluginName;
 	private final String pluginPath;
 	private final String jobType;
 	private final int order;
 	private boolean hidden;
-	
+
+	public static final Comparator<ViewerPlugin> COMPARATOR = 
+			new Comparator<ViewerPlugin>() {
+		@Override
+		public int compare(ViewerPlugin o1, ViewerPlugin o2) {
+			return o1.getOrder() - o2.getOrder();
+		}
+	};
+
 	public ViewerPlugin(
 			String pluginName, 
 			String pluginPath,