Fixed remove behavior for Device and Link Store

Change-Id: I2d6c6a48f9b92136c2f0734d0216f9f3b05b4d8c
diff --git a/core/store/dist/src/main/java/org/onlab/onos/store/device/impl/GossipDeviceStore.java b/core/store/dist/src/main/java/org/onlab/onos/store/device/impl/GossipDeviceStore.java
index 996dfa7..ac38e23 100644
--- a/core/store/dist/src/main/java/org/onlab/onos/store/device/impl/GossipDeviceStore.java
+++ b/core/store/dist/src/main/java/org/onlab/onos/store/device/impl/GossipDeviceStore.java
@@ -33,6 +33,7 @@
 import org.onlab.onos.store.ClockService;
 import org.onlab.onos.store.Timestamp;
 import org.onlab.onos.store.common.impl.Timestamped;
+import org.onlab.util.NewConcurrentHashMap;
 import org.slf4j.Logger;
 
 import java.util.ArrayList;
@@ -136,8 +137,7 @@
 
         // Collection of DeviceDescriptions for a Device
         ConcurrentMap<ProviderId, DeviceDescriptions> providerDescs
-            = createIfAbsentUnchecked(deviceDescs, deviceId,
-                    new InitConcurrentHashMap<ProviderId, DeviceDescriptions>());
+            = getDeviceDescriptions(deviceId);
 
 
         DeviceDescriptions descs
@@ -223,8 +223,7 @@
     @Override
     public DeviceEvent markOffline(DeviceId deviceId) {
         ConcurrentMap<ProviderId, DeviceDescriptions> providerDescs
-            = createIfAbsentUnchecked(deviceDescs, deviceId,
-                    new InitConcurrentHashMap<ProviderId, DeviceDescriptions>());
+            = getDeviceDescriptions(deviceId);
 
         // locking device
         synchronized (providerDescs) {
@@ -358,7 +357,13 @@
     // exist, it creates and registers a new one.
     private ConcurrentMap<PortNumber, Port> getPortMap(DeviceId deviceId) {
         return createIfAbsentUnchecked(devicePorts, deviceId,
-                new InitConcurrentHashMap<PortNumber, Port>());
+                NewConcurrentHashMap.<PortNumber, Port>ifNeeded());
+    }
+
+    private ConcurrentMap<ProviderId, DeviceDescriptions> getDeviceDescriptions(
+            DeviceId deviceId) {
+        return createIfAbsentUnchecked(deviceDescs, deviceId,
+                NewConcurrentHashMap.<ProviderId, DeviceDescriptions>ifNeeded());
     }
 
     @Override
@@ -438,10 +443,16 @@
 
     @Override
     public DeviceEvent removeDevice(DeviceId deviceId) {
-        Device device = devices.remove(deviceId);
-        // FIXME: should we be removing deviceDescs also?
-        return device == null ? null :
-            new DeviceEvent(DEVICE_REMOVED, device, null);
+        ConcurrentMap<ProviderId, DeviceDescriptions> descs = getDeviceDescriptions(deviceId);
+        synchronized (descs) {
+            Device device = devices.remove(deviceId);
+            // should DEVICE_REMOVED carry removed ports?
+            devicePorts.get(deviceId).clear();
+            availableDevices.remove(deviceId);
+            descs.clear();
+            return device == null ? null :
+                new DeviceEvent(DEVICE_REMOVED, device, null);
+        }
     }
 
     /**
@@ -546,14 +557,6 @@
         return fallBackPrimary;
     }
 
-    private static final class InitConcurrentHashMap<K, V> implements
-            ConcurrentInitializer<ConcurrentMap<K, V>> {
-        @Override
-        public ConcurrentMap<K, V> get() throws ConcurrentException {
-            return new ConcurrentHashMap<>();
-        }
-    }
-
     public static final class InitDeviceDescs
         implements ConcurrentInitializer<DeviceDescriptions> {
 
diff --git a/core/store/dist/src/test/java/org/onlab/onos/store/device/impl/GossipDeviceStoreTest.java b/core/store/dist/src/test/java/org/onlab/onos/store/device/impl/GossipDeviceStoreTest.java
index 09c3098..94de9b2 100644
--- a/core/store/dist/src/test/java/org/onlab/onos/store/device/impl/GossipDeviceStoreTest.java
+++ b/core/store/dist/src/test/java/org/onlab/onos/store/device/impl/GossipDeviceStoreTest.java
@@ -116,17 +116,19 @@
         deviceClockManager.deactivate();
     }
 
-    private void putDevice(DeviceId deviceId, String swVersion) {
+    private void putDevice(DeviceId deviceId, String swVersion,
+                           SparseAnnotations... annotations) {
         DeviceDescription description =
                 new DefaultDeviceDescription(deviceId.uri(), SWITCH, MFR,
-                        HW, swVersion, SN);
+                        HW, swVersion, SN, annotations);
         deviceStore.createOrUpdateDevice(PID, deviceId, description);
     }
 
-    private void putDeviceAncillary(DeviceId deviceId, String swVersion) {
+    private void putDeviceAncillary(DeviceId deviceId, String swVersion,
+                                    SparseAnnotations... annotations) {
         DeviceDescription description =
                 new DefaultDeviceDescription(deviceId.uri(), SWITCH, MFR,
-                        HW, swVersion, SN);
+                        HW, swVersion, SN, annotations);
         deviceStore.createOrUpdateDevice(PIDA, deviceId, description);
     }
 
@@ -448,16 +450,37 @@
 
     @Test
     public final void testRemoveDevice() {
-        putDevice(DID1, SW1);
+        putDevice(DID1, SW1, A1);
+        List<PortDescription> pds = Arrays.<PortDescription>asList(
+                new DefaultPortDescription(P1, true, A2)
+                );
+        deviceStore.updatePorts(PID, DID1, pds);
         putDevice(DID2, SW1);
 
         assertEquals(2, deviceStore.getDeviceCount());
+        assertEquals(1, deviceStore.getPorts(DID1).size());
+        assertAnnotationsEquals(deviceStore.getDevice(DID1).annotations(), A1);
+        assertAnnotationsEquals(deviceStore.getPort(DID1, P1).annotations(), A2);
 
         DeviceEvent event = deviceStore.removeDevice(DID1);
         assertEquals(DEVICE_REMOVED, event.type());
         assertDevice(DID1, SW1, event.subject());
 
         assertEquals(1, deviceStore.getDeviceCount());
+        assertEquals(0, deviceStore.getPorts(DID1).size());
+
+        // putBack Device, Port w/o annotation
+        putDevice(DID1, SW1);
+        List<PortDescription> pds2 = Arrays.<PortDescription>asList(
+                new DefaultPortDescription(P1, true)
+                );
+        deviceStore.updatePorts(PID, DID1, pds2);
+
+        // annotations should not survive
+        assertEquals(2, deviceStore.getDeviceCount());
+        assertEquals(1, deviceStore.getPorts(DID1).size());
+        assertAnnotationsEquals(deviceStore.getDevice(DID1).annotations());
+        assertAnnotationsEquals(deviceStore.getPort(DID1, P1).annotations());
     }
 
     // If Delegates should be called only on remote events,
diff --git a/core/store/trivial/src/main/java/org/onlab/onos/store/trivial/impl/SimpleDeviceStore.java b/core/store/trivial/src/main/java/org/onlab/onos/store/trivial/impl/SimpleDeviceStore.java
index 4eda2fc..926c8a7 100644
--- a/core/store/trivial/src/main/java/org/onlab/onos/store/trivial/impl/SimpleDeviceStore.java
+++ b/core/store/trivial/src/main/java/org/onlab/onos/store/trivial/impl/SimpleDeviceStore.java
@@ -53,7 +53,6 @@
 import static org.apache.commons.lang3.concurrent.ConcurrentUtils.createIfAbsentUnchecked;
 import static org.onlab.onos.net.DefaultAnnotations.merge;
 
-// TODO: synchronization should be done in more fine-grained manner.
 /**
  * Manages inventory of infrastructure devices using trivial in-memory
  * structures implementation.
@@ -329,11 +328,18 @@
 
     @Override
     public DeviceEvent removeDevice(DeviceId deviceId) {
-        synchronized (this) {
+        ConcurrentMap<ProviderId, DeviceDescriptions> descs = getDeviceDescriptions(deviceId);
+        synchronized (descs) {
             Device device = devices.remove(deviceId);
-            // FIXME: should we be removing deviceDescs also?
+            // should DEVICE_REMOVED carry removed ports?
+            ConcurrentMap<PortNumber, Port> ports = devicePorts.get(deviceId);
+            if (ports != null) {
+                ports.clear();
+            }
+            availableDevices.remove(deviceId);
+            descs.clear();
             return device == null ? null :
-                    new DeviceEvent(DEVICE_REMOVED, device, null);
+                new DeviceEvent(DEVICE_REMOVED, device, null);
         }
     }
 
diff --git a/core/store/trivial/src/main/java/org/onlab/onos/store/trivial/impl/SimpleLinkStore.java b/core/store/trivial/src/main/java/org/onlab/onos/store/trivial/impl/SimpleLinkStore.java
index 230683d..1578384 100644
--- a/core/store/trivial/src/main/java/org/onlab/onos/store/trivial/impl/SimpleLinkStore.java
+++ b/core/store/trivial/src/main/java/org/onlab/onos/store/trivial/impl/SimpleLinkStore.java
@@ -47,7 +47,6 @@
 import static com.google.common.collect.Multimaps.synchronizedSetMultimap;
 import static com.google.common.base.Predicates.notNull;
 
-// TODO: Add support for multiple provider and annotations
 /**
  * Manages inventory of infrastructure links using trivial in-memory structures
  * implementation.
@@ -230,7 +229,7 @@
         ConcurrentMap<ProviderId, LinkDescription> descs = getLinkDescriptions(key);
         synchronized (descs) {
             Link link = links.remove(key);
-            // FIXME: should we be removing deviceDescs also?
+            descs.clear();
             if (link != null) {
                 srcLinks.remove(link.src().deviceId(), key);
                 dstLinks.remove(link.dst().deviceId(), key);
diff --git a/core/store/trivial/src/test/java/org/onlab/onos/store/trivial/impl/SimpleDeviceStoreTest.java b/core/store/trivial/src/test/java/org/onlab/onos/store/trivial/impl/SimpleDeviceStoreTest.java
index 34bb1f4..146086a 100644
--- a/core/store/trivial/src/test/java/org/onlab/onos/store/trivial/impl/SimpleDeviceStoreTest.java
+++ b/core/store/trivial/src/test/java/org/onlab/onos/store/trivial/impl/SimpleDeviceStoreTest.java
@@ -103,17 +103,19 @@
         simpleDeviceStore.deactivate();
     }
 
-    private void putDevice(DeviceId deviceId, String swVersion) {
+    private void putDevice(DeviceId deviceId, String swVersion,
+                           SparseAnnotations... annotations) {
         DeviceDescription description =
                 new DefaultDeviceDescription(deviceId.uri(), SWITCH, MFR,
-                        HW, swVersion, SN);
+                        HW, swVersion, SN, annotations);
         deviceStore.createOrUpdateDevice(PID, deviceId, description);
     }
 
-    private void putDeviceAncillary(DeviceId deviceId, String swVersion) {
+    private void putDeviceAncillary(DeviceId deviceId, String swVersion,
+                                    SparseAnnotations... annotations) {
         DeviceDescription description =
                 new DefaultDeviceDescription(deviceId.uri(), SWITCH, MFR,
-                        HW, swVersion, SN);
+                        HW, swVersion, SN, annotations);
         deviceStore.createOrUpdateDevice(PIDA, deviceId, description);
     }
 
@@ -437,16 +439,37 @@
 
     @Test
     public final void testRemoveDevice() {
-        putDevice(DID1, SW1);
+        putDevice(DID1, SW1, A1);
+        List<PortDescription> pds = Arrays.<PortDescription>asList(
+                new DefaultPortDescription(P1, true, A2)
+                );
+        deviceStore.updatePorts(PID, DID1, pds);
         putDevice(DID2, SW1);
 
         assertEquals(2, deviceStore.getDeviceCount());
+        assertEquals(1, deviceStore.getPorts(DID1).size());
+        assertAnnotationsEquals(deviceStore.getDevice(DID1).annotations(), A1);
+        assertAnnotationsEquals(deviceStore.getPort(DID1, P1).annotations(), A2);
 
         DeviceEvent event = deviceStore.removeDevice(DID1);
         assertEquals(DEVICE_REMOVED, event.type());
         assertDevice(DID1, SW1, event.subject());
 
         assertEquals(1, deviceStore.getDeviceCount());
+        assertEquals(0, deviceStore.getPorts(DID1).size());
+
+        // putBack Device, Port w/o annotation
+        putDevice(DID1, SW1);
+        List<PortDescription> pds2 = Arrays.<PortDescription>asList(
+                new DefaultPortDescription(P1, true)
+                );
+        deviceStore.updatePorts(PID, DID1, pds2);
+
+        // annotations should not survive
+        assertEquals(2, deviceStore.getDeviceCount());
+        assertEquals(1, deviceStore.getPorts(DID1).size());
+        assertAnnotationsEquals(deviceStore.getDevice(DID1).annotations());
+        assertAnnotationsEquals(deviceStore.getPort(DID1, P1).annotations());
     }
 
     // If Delegates should be called only on remote events,
diff --git a/core/store/trivial/src/test/java/org/onlab/onos/store/trivial/impl/SimpleLinkStoreTest.java b/core/store/trivial/src/test/java/org/onlab/onos/store/trivial/impl/SimpleLinkStoreTest.java
index 47da868..8a16609 100644
--- a/core/store/trivial/src/test/java/org/onlab/onos/store/trivial/impl/SimpleLinkStoreTest.java
+++ b/core/store/trivial/src/test/java/org/onlab/onos/store/trivial/impl/SimpleLinkStoreTest.java
@@ -91,16 +91,17 @@
     }
 
     private void putLink(DeviceId srcId, PortNumber srcNum,
-                         DeviceId dstId, PortNumber dstNum, Type type) {
+                         DeviceId dstId, PortNumber dstNum, Type type,
+                         SparseAnnotations... annotations) {
         ConnectPoint src = new ConnectPoint(srcId, srcNum);
         ConnectPoint dst = new ConnectPoint(dstId, dstNum);
-        linkStore.createOrUpdateLink(PID, new DefaultLinkDescription(src, dst, type));
+        linkStore.createOrUpdateLink(PID, new DefaultLinkDescription(src, dst, type, annotations));
     }
 
-    private void putLink(LinkKey key, Type type) {
+    private void putLink(LinkKey key, Type type, SparseAnnotations... annotations) {
         putLink(key.src().deviceId(), key.src().port(),
                 key.dst().deviceId(), key.dst().port(),
-                type);
+                type, annotations);
     }
 
     private static void assertLink(DeviceId srcId, PortNumber srcNum,
@@ -351,8 +352,8 @@
         LinkKey linkId1 = new LinkKey(d1P1, d2P2);
         LinkKey linkId2 = new LinkKey(d2P2, d1P1);
 
-        putLink(linkId1, DIRECT);
-        putLink(linkId2, DIRECT);
+        putLink(linkId1, DIRECT, A1);
+        putLink(linkId2, DIRECT, A2);
 
         // DID1,P1 => DID2,P2
         // DID2,P2 => DID1,P1
@@ -360,10 +361,17 @@
 
         LinkEvent event = linkStore.removeLink(d1P1, d2P2);
         assertEquals(LINK_REMOVED, event.type());
+        assertAnnotationsEquals(event.subject().annotations(), A1);
         LinkEvent event2 = linkStore.removeLink(d1P1, d2P2);
         assertNull(event2);
 
         assertLink(linkId2, DIRECT, linkStore.getLink(d2P2, d1P1));
+        assertAnnotationsEquals(linkStore.getLink(d2P2, d1P1).annotations(), A2);
+
+        // annotations, etc. should not survive remove
+        putLink(linkId1, DIRECT);
+        assertLink(linkId1, DIRECT, linkStore.getLink(d1P1, d2P2));
+        assertAnnotationsEquals(linkStore.getLink(d1P1, d2P2).annotations());
     }
 
     @Test