/*
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You under the Apache License, Version 2.0
 * (the "License"); you may not use this file except in compliance with
 * the License.  You may obtain a copy of the License at
 *
 *      http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
package org.apache.wicket.protocol.ws.api.registry;

import java.util.ArrayList;
import java.util.Collection;
import java.util.concurrent.ConcurrentMap;

import org.apache.wicket.Application;
import org.apache.wicket.MetaDataKey;
import org.apache.wicket.protocol.ws.api.IWebSocketConnection;
import org.apache.wicket.util.lang.Args;
import org.apache.wicket.util.lang.Generics;

/**
 * A registry that keeps all currently opened web socket connections in
 * maps in Application's meta data.
 *
 * TODO remove the synchronizations below and use ConcurrentMap#putIfAbsent()
 *
 * @since 6.0
 */
public class SimpleWebSocketConnectionRegistry implements IWebSocketConnectionRegistry
{
	private static final MetaDataKey<ConcurrentMap<String, ConcurrentMap<IKey, IWebSocketConnection>>> KEY =
			new MetaDataKey<ConcurrentMap<String, ConcurrentMap<IKey, IWebSocketConnection>>>()
	{
	};

	@Override
	public IWebSocketConnection getConnection(Application application, String sessionId, IKey key)
	{
		Args.notNull(application, "application");
		Args.notNull(sessionId, "sessionId");
		Args.notNull(key, "key");

		IWebSocketConnection connection = null;
		ConcurrentMap<String, ConcurrentMap<IKey, IWebSocketConnection>> connectionsBySession = application.getMetaData(KEY);
		if (connectionsBySession != null)
		{
			ConcurrentMap<IKey, IWebSocketConnection> connectionsByPage = connectionsBySession.get(sessionId);
			if (connectionsByPage != null)
			{
				connection = connectionsByPage.get(key);
			}
		}
		return connection;
	}

	/**
	 * Returns a collection of currently active websockets. The connections might close at any time.
	 *
	 * @param application
	 * @return
	 */
	public Collection<IWebSocketConnection> getConnections(Application application)
	{
		Args.notNull(application, "application");

		Collection<IWebSocketConnection> connections = new ArrayList<>();
		ConcurrentMap<String, ConcurrentMap<IKey, IWebSocketConnection>> connectionsBySession = application.getMetaData(KEY);
		if (connectionsBySession != null)
		{
			for (ConcurrentMap<IKey, IWebSocketConnection> connectionsByPage : connectionsBySession.values())
			{
				connections.addAll(connectionsByPage.values());
			}
		}
		return connections;
	}

	@Override
	public void setConnection(Application application, String sessionId, IKey key, IWebSocketConnection connection)
	{
		Args.notNull(application, "application");
		Args.notNull(sessionId, "sessionId");
		Args.notNull(key, "key");

		ConcurrentMap<String, ConcurrentMap<IKey, IWebSocketConnection>> connectionsBySession = application.getMetaData(KEY);
		if (connectionsBySession == null)
		{
			synchronized (KEY)
			{
				connectionsBySession = application.getMetaData(KEY);
				if (connectionsBySession == null)
				{
					connectionsBySession = Generics.newConcurrentHashMap();
					application.setMetaData(KEY, connectionsBySession);
				}
			}
		}

		ConcurrentMap<IKey, IWebSocketConnection> connectionsByPage = connectionsBySession.get(sessionId);
		if (connectionsByPage == null && connection != null)
		{
			synchronized (connectionsBySession)
			{
				connectionsByPage = connectionsBySession.get(sessionId);
				if (connectionsByPage == null)
				{
					connectionsByPage = Generics.newConcurrentHashMap();
					connectionsBySession.put(sessionId, connectionsByPage);
				}
			}
		}

		if (connection != null)
		{
			connectionsByPage.put(key, connection);
		}
		else if (connectionsByPage != null)
		{
			connectionsByPage.remove(key);
			if (connectionsByPage.isEmpty())
			{
				connectionsBySession.remove(sessionId);
			}
		}
	}

	@Override
	public void removeConnection(Application application, String sessionId, IKey key)
	{
		setConnection(application, sessionId, key, null);
	}
}

@@ -2463,15 +2463,88 @@
 %09 * 
-@return
+         The application%0A%09 * @return a collection of currently active websockets
 %0A%09 *
@@ -3921,53 +3921,8 @@
 %09%09%7B%0A
-%09%09%09synchronized (connectionsBySession)%0A%09%09%09%7B%0A%09
 %09%09%09c
@@ -3977,17 +3977,16 @@
 ionId);%0A
-%09
 %09%09%09if (c
@@ -4006,36 +4006,34 @@
 age == null)%0A%09%09%09
-%09
 %7B%0A
-%09
 %09%09%09%09connectionsB
@@ -4073,25 +4073,72 @@
 hMap();%0A%09%09%09%09
-%09
+ConcurrentMap%3CIKey, IWebSocketConnection%3E old = 
 connectionsB
@@ -4149,16 +4149,24 @@
 sion.put
+IfAbsent
 (session
@@ -4189,16 +4189,73 @@
 yPage);%0A
+%09%09%09%09if (old != null)%0A%09%09%09%09%7B%0A%09%09%09%09%09connectionsByPage = old;%0A
 %09%09%09%09%7D%0A%09%09

Improve usage of concurrent map - no need of synchronization.