From 4b7a6e26fa7cb8fac5c6781068dc8fbfcb1d5033 Mon Sep 17 00:00:00 2001
From: Simon Bernard <sbernard@sierrawireless.com>
Date: Thu, 31 May 2018 17:16:09 +0200
Subject: [PATCH] Add registrationId index in InMemoryRegistrationStore

---
 .../impl/InMemoryRegistrationStore.java         | 17 ++++++++---------
 1 file changed, 8 insertions(+), 9 deletions(-)

diff --git a/leshan-server-cf/src/main/java/org/eclipse/leshan/server/californium/impl/InMemoryRegistrationStore.java b/leshan-server-cf/src/main/java/org/eclipse/leshan/server/californium/impl/InMemoryRegistrationStore.java
index 941fe7533c..4c9169762d 100644
--- a/leshan-server-cf/src/main/java/org/eclipse/leshan/server/californium/impl/InMemoryRegistrationStore.java
+++ b/leshan-server-cf/src/main/java/org/eclipse/leshan/server/californium/impl/InMemoryRegistrationStore.java
@@ -67,6 +67,7 @@ public class InMemoryRegistrationStore implements CaliforniumRegistrationStore,
     // Data structure
     private final Map<String /* end-point */, Registration> regsByEp = new HashMap<>();
     private final Map<InetSocketAddress, Registration> regsByAddr = new HashMap<>();
+    private final Map<String /* reg-id */, Registration> regsByRegId = new HashMap<>();
     private Map<Token, org.eclipse.californium.core.observe.Observation> obsByToken = new HashMap<>();
     private Map<String, Set<Token>> tokensByRegId = new HashMap<>();
 
@@ -101,12 +102,14 @@ public Deregistration addRegistration(Registration registration) {
             lock.writeLock().lock();
 
             Registration registrationRemoved = regsByEp.put(registration.getEndpoint(), registration);
+            regsByRegId.put(registration.getId(), registration);
             // If a registration is already associated to this address we don't care as we only want to keep the most
             // recent binding.
             regsByAddr.put(registration.getSocketAddress(), registration);
             if (registrationRemoved != null) {
                 Collection<Observation> observationsRemoved = unsafeRemoveAllObservations(registrationRemoved.getId());
                 removeFromMap(regsByAddr, registrationRemoved.getSocketAddress(), registrationRemoved);
+                removeFromMap(regsByRegId, registrationRemoved.getId(), registrationRemoved);
                 return new Deregistration(registrationRemoved, observationsRemoved);
             }
         } finally {
@@ -130,6 +133,9 @@ public UpdatedRegistration updateRegistration(RegistrationUpdate update) {
                 // recent binding.
                 regsByAddr.put(updatedRegistration.getSocketAddress(), updatedRegistration);
                 removeFromMap(regsByAddr, registration.getSocketAddress(), registration);
+
+                regsByRegId.put(updatedRegistration.getId(), updatedRegistration);
+
                 return new UpdatedRegistration(registration, updatedRegistration);
             }
         } finally {
@@ -141,15 +147,7 @@ public UpdatedRegistration updateRegistration(RegistrationUpdate update) {
     public Registration getRegistration(String registrationId) {
         try {
             lock.readLock().lock();
-            // TODO we should create an index instead of iterate all over the collection
-            if (registrationId != null) {
-                for (Registration registration : regsByEp.values()) {
-                    if (registrationId.equals(registration.getId())) {
-                        return registration;
-                    }
-                }
-            }
-            return null;
+            return regsByRegId.get(registrationId);
         } finally {
             lock.readLock().unlock();
         }
@@ -195,6 +193,7 @@ public Deregistration removeRegistration(String registrationId) {
                 Collection<Observation> observationsRemoved = unsafeRemoveAllObservations(registration.getId());
                 regsByEp.remove(registration.getEndpoint());
                 removeFromMap(regsByAddr, registration.getSocketAddress(), registration);
+                removeFromMap(regsByRegId, registration.getId(), registration);
                 return new Deregistration(registration, observationsRemoved);
             }
             return null;