001/*
002 * #%L
003 * Netarchivesuite - harvester
004 * %%
005 * Copyright (C) 2005 - 2014 The Royal Danish Library, the Danish State and University Library,
006 *             the National Library of France and the Austrian National Library.
007 * %%
008 * This program is free software: you can redistribute it and/or modify
009 * it under the terms of the GNU Lesser General Public License as
010 * published by the Free Software Foundation, either version 2.1 of the
011 * License, or (at your option) any later version.
012 * 
013 * This program is distributed in the hope that it will be useful,
014 * but WITHOUT ANY WARRANTY; without even the implied warranty of
015 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
016 * GNU General Lesser Public License for more details.
017 * 
018 * You should have received a copy of the GNU General Lesser Public
019 * License along with this program.  If not, see
020 * <http://www.gnu.org/licenses/lgpl-2.1.html>.
021 * #L%
022 */
023package dk.netarkivet.harvester.scheduler;
024
025import java.util.ArrayList;
026import java.util.HashMap;
027import java.util.Map;
028import java.util.Set;
029import java.util.TreeSet;
030
031import org.slf4j.Logger;
032import org.slf4j.LoggerFactory;
033
034import dk.netarkivet.harvester.datamodel.HarvestChannel;
035import dk.netarkivet.harvester.harvesting.distribute.HarvesterRegistrationRequest;
036
037/**
038 * Keeps track of the number of {@link HarvesterRegistrationRequest}s that have been received per channel, which allows
039 * to know if a HarvestControllers are registered to a given {@link HarvestChannel}.
040 */
041public class HarvestChannelRegistry {
042
043    /** The class logger. */
044    private static final Logger LOG = LoggerFactory.getLogger(HarvestChannelRegistry.class);
045
046    private Map<String, Set<String>> harvesterChannelRegistry = new HashMap<String, Set<String>>();
047
048    public synchronized void register(final String channelName, final String harvesterInstanceId) {
049        // First remove any reference to this instanceId
050        // This is done in case a Harvester has been unexpectedly shut down and restarted
051        clearHarvester(harvesterInstanceId);
052
053        Set<String> instanceIds = harvesterChannelRegistry.get(channelName);
054        if (instanceIds == null) {
055            instanceIds = new TreeSet<String>();
056        }
057        instanceIds.add(harvesterInstanceId);
058        harvesterChannelRegistry.put(channelName, instanceIds);
059
060        LOG.info("Harvester {} registered on channel {}", harvesterInstanceId, channelName);
061        logStatus();
062    }
063
064    public synchronized boolean isRegistered(final String channelName) {
065        return harvesterChannelRegistry.containsKey(channelName);
066    }
067    
068    public synchronized boolean isRegisteredToChannel(final String harvesterInstanceId, final String channelName) {
069        if (!isRegistered(channelName)){
070                return false;
071        } else {
072                return harvesterChannelRegistry.get(channelName).contains(harvesterInstanceId);
073        }
074    }
075
076    private void logStatus() {
077        if (LOG.isInfoEnabled()) {
078            String msg = HarvestChannelRegistry.class.getSimpleName() + " status:";
079            for (String channel : harvesterChannelRegistry.keySet()) {
080                msg += "\n\t- " + channel + " { ";
081                for (String id : harvesterChannelRegistry.get(channel)) {
082                    msg += id + ", ";
083                }
084                msg = msg.substring(0, msg.lastIndexOf(",")) + " }";
085            }
086            LOG.info(msg);
087        }
088    }
089
090    /**
091     * Clears any registration data for a given harvester instance id.
092     *
093     * @param harvesterInstanceId a harvester instance id
094     */
095    private void clearHarvester(final String harvesterInstanceId) {
096        ArrayList<String> keysToRemove = new ArrayList<String>();
097        for (String channel : harvesterChannelRegistry.keySet()) {
098            Set<String> instanceIds = harvesterChannelRegistry.get(channel);
099            if (instanceIds.contains(harvesterInstanceId)) {
100                instanceIds.remove(harvesterInstanceId);
101                if (instanceIds.isEmpty()) {
102                    keysToRemove.add(channel);
103                }
104                LOG.info("Cleared former registration of '{}' to '{}'", channel, harvesterInstanceId);
105            }
106        }
107        for (String key : keysToRemove) {
108            harvesterChannelRegistry.remove(key);
109        }
110    }
111
112}