// Copyright 2019 Oath Inc. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.

package com.yahoo.vespa.model;

import static com.yahoo.config.provision.NetworkPorts.Allocation;

import java.util.Collection;
import java.util.HashMap;
import java.util.Map;
import java.util.TreeMap;

public class PortFinder {

    private final Map<String, Allocation> byKeys = new HashMap<>();
    private final Map<Integer, Allocation> byPorts = new TreeMap<>();

    /** force add the given allocation, removing any conflicting ones */
    public void use(Allocation allocation) {
        String key = allocation.key();
        if (byKeys.containsKey(key)) {
            if (byKeys.get(key).port == allocation.port) {
                return; // already OK
            }
            Allocation toRemove = byKeys.remove(key);
            byPorts.remove(toRemove.port);
        }
        if (byPorts.containsKey(allocation.port)) {
            Allocation toRemove = byPorts.remove(allocation.port);
            byKeys.remove(toRemove.key());
        }
        byPorts.put(allocation.port, allocation);
        byKeys.put(key, allocation);
    }

    public int findPort(Allocation request) {
        String key = request.key();
        if (byKeys.containsKey(key)) {
            return byKeys.get(key).port;
        }
        int port = request.port;
        while (byPorts.containsKey(port)) {
            ++port;
        }
        return port;
    }

    public PortFinder(Collection<Allocation> allocations) {
        for (Allocation a : allocations) {
            use(a);
        }
    }

    public Collection<Allocation> allocations() {
        return byPorts.values();
    }

}

@@ -314,33 +314,184 @@
 p;%0A%0A
-public class PortFinder %7B
+import java.util.logging.Logger;%0Aimport java.util.logging.Level;%0A%0Apublic class PortFinder %7B%0A%0A    private static final Logger log = Logger.getLogger(PortFinder.class.getName());
 %0A%0A  
@@ -1440,22 +1440,26 @@
         
-return
+int port =
  byKeys.
@@ -1469,24 +1469,132 @@
 (key).port;%0A
+            log.log(Level.INFO, %22Re-using port %22+port+%22 for allocation %22+request);%0A            return port;%0A
         %7D%0A  

add info-level logging so we can verify that port re-use works as expected