From 7f357d68cca8930235683b43460e56973efdb6a5 Mon Sep 17 00:00:00 2001
From: Chris Nauroth <cnauroth@apache.org>
Date: Wed, 16 Jul 2014 18:50:10 +0000
Subject: [PATCH] HADOOP-10839. Add unregisterSource() to MetricsSystem API.
 Contributed by Shanyu Zhao.

git-svn-id: https://svn.apache.org/repos/asf/hadoop/common/trunk@1611134 13f79535-47bb-0310-9956-ffa450edef68
---
 hadoop-common-project/hadoop-common/CHANGES.txt |  3 +++
 .../apache/hadoop/metrics2/MetricsSystem.java   |  6 ++++++
 .../hadoop/metrics2/impl/MetricsSystemImpl.java | 11 +++++++++++
 .../metrics2/impl/TestMetricsSystemImpl.java    | 17 +++++++++++++++++
 4 files changed, 37 insertions(+)

diff --git a/hadoop-common-project/hadoop-common/CHANGES.txt b/hadoop-common-project/hadoop-common/CHANGES.txt
index d52558620d..75dba1892a 100644
--- a/hadoop-common-project/hadoop-common/CHANGES.txt
+++ b/hadoop-common-project/hadoop-common/CHANGES.txt
@@ -409,6 +409,9 @@ Release 2.6.0 - UNRELEASED
     HADOOP-10845. Add common tests for ACLs in combination with viewfs.
     (Stephen Chu via cnauroth)
 
+    HADOOP-10839. Add unregisterSource() to MetricsSystem API.
+    (Shanyu Zhao via cnauroth)
+
   OPTIMIZATIONS
 
   BUG FIXES
diff --git a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/metrics2/MetricsSystem.java b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/metrics2/MetricsSystem.java
index e853319c4e..a277abd6e1 100644
--- a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/metrics2/MetricsSystem.java
+++ b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/metrics2/MetricsSystem.java
@@ -54,6 +54,12 @@ public abstract class MetricsSystem implements MetricsSystemMXBean {
    */
   public abstract <T> T register(String name, String desc, T source);
 
+  /**
+   * Unregister a metrics source
+   * @param name of the source. This is the name you use to call register()
+   */
+  public abstract void unregisterSource(String name);
+
   /**
    * Register a metrics source (deriving name and description from the object)
    * @param <T>   the actual type of the source object
diff --git a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/metrics2/impl/MetricsSystemImpl.java b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/metrics2/impl/MetricsSystemImpl.java
index cf2dda4e38..722abd95c4 100644
--- a/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/metrics2/impl/MetricsSystemImpl.java
+++ b/hadoop-common-project/hadoop-common/src/main/java/org/apache/hadoop/metrics2/impl/MetricsSystemImpl.java
@@ -232,6 +232,17 @@ T register(String name, String desc, T source) {
     return source;
   }
 
+  @Override public synchronized
+  void unregisterSource(String name) {
+    if (sources.containsKey(name)) {
+      sources.get(name).stop();
+      sources.remove(name);
+    }
+    if (allSources.containsKey(name)) {
+      allSources.remove(name);
+    }
+  }
+
   synchronized
   void registerSource(String name, String desc, MetricsSource source) {
     checkNotNull(config, "config");
diff --git a/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/metrics2/impl/TestMetricsSystemImpl.java b/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/metrics2/impl/TestMetricsSystemImpl.java
index 63dcbb2afa..0122045d38 100644
--- a/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/metrics2/impl/TestMetricsSystemImpl.java
+++ b/hadoop-common-project/hadoop-common/src/test/java/org/apache/hadoop/metrics2/impl/TestMetricsSystemImpl.java
@@ -380,6 +380,23 @@ public void flush() {
     ms.shutdown();
   }
 
+  @Test public void testUnregisterSource() {
+    MetricsSystem ms = new MetricsSystemImpl();
+    TestSource ts1 = new TestSource("ts1");
+    TestSource ts2 = new TestSource("ts2");
+    ms.register("ts1", "", ts1);
+    ms.register("ts2", "", ts2);
+    MetricsSource s1 = ms.getSource("ts1");
+    assertNotNull(s1);
+    // should work when metrics system is not started
+    ms.unregisterSource("ts1");
+    s1 = ms.getSource("ts1");
+    assertNull(s1);
+    MetricsSource s2 = ms.getSource("ts2");
+    assertNotNull(s2);
+    ms.shutdown();
+  }
+
   private void checkMetricsRecords(List<MetricsRecord> recs) {
     LOG.debug(recs);
     MetricsRecord r = recs.get(0);