org.wildfly.clustering.server.infinispan.registry.CacheRegistry Maven / Gradle / Ivy

Go to download
/*
 * Copyright The WildFly Authors
 * SPDX-License-Identifier: Apache-2.0
 */
package org.wildfly.clustering.server.infinispan.registry;

import java.util.Collections;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.Executor;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;
import java.util.function.Supplier;
import java.util.stream.Collectors;

import org.infinispan.Cache;
import org.infinispan.commons.CacheException;
import org.infinispan.context.Flag;
import org.infinispan.distribution.ch.ConsistentHash;
import org.infinispan.manager.EmbeddedCacheManager;
import org.infinispan.notifications.Listener;
import org.infinispan.notifications.Listener.Observation;
import org.infinispan.notifications.cachelistener.annotation.CacheEntryCreated;
import org.infinispan.notifications.cachelistener.annotation.CacheEntryModified;
import org.infinispan.notifications.cachelistener.annotation.CacheEntryRemoved;
import org.infinispan.notifications.cachelistener.annotation.TopologyChanged;
import org.infinispan.notifications.cachelistener.event.CacheEntryEvent;
import org.infinispan.notifications.cachelistener.event.CacheEntryRemovedEvent;
import org.infinispan.notifications.cachelistener.event.Event;
import org.infinispan.notifications.cachelistener.event.TopologyChangedEvent;
import org.infinispan.remoting.transport.Address;
import org.jboss.logging.Logger;
import org.wildfly.clustering.cache.batch.Batch;
import org.wildfly.clustering.cache.infinispan.embedded.distribution.Locality;
import org.wildfly.clustering.cache.infinispan.embedded.listener.KeyFilter;
import org.wildfly.clustering.context.DefaultExecutorService;
import org.wildfly.clustering.context.ExecutorServiceFactory;
import org.wildfly.clustering.server.Registration;
import org.wildfly.clustering.server.infinispan.CacheContainerGroup;
import org.wildfly.clustering.server.infinispan.CacheContainerGroupMember;
import org.wildfly.clustering.server.registry.Registry;
import org.wildfly.clustering.server.registry.RegistryListener;

/**
 * Clustered {@link Registry} backed by an Infinispan cache.
 * @author Paul Ferraro
 * @param  key type
 * @param  value type
 */
@Listener(observation = Observation.POST)
public class CacheRegistry implements CacheContainerRegistry {
	private static final Logger LOGGER = Logger.getLogger(CacheRegistry.class);

	private final Map, ExecutorService> listeners = new ConcurrentHashMap<>();
	private final Cache> cache;
	private final Supplier batchFactory;
	private final CacheContainerGroup group;
	private final Runnable closeTask;
	private final Map.Entry entry;
	private final Executor executor;
	private final Function, ExecutorService> executorServiceFactory = new Function<>() {
		@Override
		public ExecutorService apply(RegistryListener listener) {
			return new DefaultExecutorService(ExecutorServiceFactory.SINGLE_THREAD, Thread.currentThread().getContextClassLoader());
		}
	};

	public CacheRegistry(CacheRegistryConfiguration config, Map.Entry entry, Runnable closeTask) {
		this.cache = config.getWriteOnlyCache();
		this.batchFactory = config.getBatchFactory();
		this.group = config.getGroup();
		this.closeTask = closeTask;
		this.executor = config.getExecutor();
		this.entry = entry;
		Address localAddress = this.cache.getCacheManager().getAddress();
		try (Batch batch = this.batchFactory.get()) {
			this.cache.put(localAddress, this.entry);
		}
		if (!this.group.isSingleton()) {
			this.cache.addListener(this, new KeyFilter<>(Address.class), null);
		}
	}

	@Override
	public void close() {
		if (!this.group.isSingleton()) {
			this.cache.removeListener(this);
		}
		Address localAddress = this.cache.getCacheManager().getAddress();
		try (Batch batch = this.batchFactory.get()) {
			// If this remove fails, the entry will be auto-removed on topology change by the new primary owner
			this.cache.getAdvancedCache().withFlags(Flag.IGNORE_RETURN_VALUES, Flag.FAIL_SILENTLY).remove(localAddress);
		} catch (CacheException e) {
			LOGGER.warn(e.getLocalizedMessage(), e);
		} finally {
			// Cleanup any unregistered listeners
			for (ExecutorService executor : this.listeners.values()) {
				this.shutdown(executor);
			}
			this.listeners.clear();
			this.closeTask.run();
		}
	}

	@Override
	public Registration register(RegistryListener listener) {
		if (this.group.isSingleton()) {
			return Registration.EMPTY;
		}
		this.listeners.computeIfAbsent(listener, this.executorServiceFactory);
		return () -> this.unregister(listener);
	}

	private void unregister(RegistryListener listener) {
		ExecutorService executor = this.listeners.remove(listener);
		if (executor != null) {
			this.shutdown(executor);
		}
	}

	@Override
	public CacheContainerGroup getGroup() {
		return this.group;
	}

	@Override
	public Map getEntries() {
		Set addresses = this.group.getMembership().getMembers().stream().map(CacheContainerGroupMember::getAddress).collect(Collectors.toUnmodifiableSet());
		Map result = new HashMap<>();
		for (Map.Entry entry : this.cache.getAdvancedCache().getAll(addresses).values()) {
			result.put(entry.getKey(), entry.getValue());
		}
		return Collections.unmodifiableMap(result);
	}

	@Override
	public Map.Entry getEntry(CacheContainerGroupMember member) {
		return this.cache.get(member.getAddress());
	}

	@TopologyChanged
	public CompletionStage topologyChanged(TopologyChangedEvent> event) {
		ConsistentHash previousHash = event.getWriteConsistentHashAtStart();
		List previousMembers = previousHash.getMembers();
		ConsistentHash hash = event.getWriteConsistentHashAtEnd();
		List members = hash.getMembers();

		if (!members.equals(previousMembers)) {
			Cache> cache = event.getCache().getAdvancedCache().withFlags(Flag.FORCE_SYNCHRONOUS);
			EmbeddedCacheManager container = cache.getCacheManager();
			Address localAddress = container.getAddress();

			// Determine which group members have left the cache view
			Set leftMembers = new HashSet<>(previousMembers);
			leftMembers.removeAll(members);

			if (!leftMembers.isEmpty()) {
				Locality locality = Locality.forConsistentHash(cache, hash);
				// We're only interested in the entries for which we are the primary owner
				Iterator addresses = leftMembers.iterator();
				while (addresses.hasNext()) {
					if (!locality.isLocal(addresses.next())) {
						addresses.remove();
					}
				}
			}

			// If this is a merge after cluster split: re-populate the cache registry with lost registry entries
			boolean restoreLocalEntry = !previousMembers.contains(localAddress);

			if (!leftMembers.isEmpty() || restoreLocalEntry) {
				try {
					this.executor.execute(() -> {
						if (!leftMembers.isEmpty()) {
							Map removed = new HashMap<>();
							try {
								for (Address leftMember: leftMembers) {
									Map.Entry old = cache.remove(leftMember);
									if (old != null) {
										removed.put(old.getKey(), old.getValue());
									}
								}
							} catch (CacheException e) {
								LOGGER.warn(e.getLocalizedMessage(), e);
							}
							if (!removed.isEmpty()) {
								this.notifyListeners(Event.Type.CACHE_ENTRY_REMOVED, removed);
							}
						}
						if (restoreLocalEntry) {
							// If we were not a member at merge start, its mapping may have been lost and need to be recreated
							try {
								if (cache.put(localAddress, this.entry) == null) {
									// Local cache events do not trigger notifications
									this.notifyListeners(Event.Type.CACHE_ENTRY_CREATED, this.entry);
								}
							} catch (CacheException e) {
								LOGGER.warn(e.getLocalizedMessage(), e);
							}
						}
					});
				} catch (RejectedExecutionException e) {
					// Executor was shutdown
				}
			}
		}
		return CompletableFuture.completedStage(null);
	}

	@CacheEntryCreated
	@CacheEntryModified
	public CompletionStage event(CacheEntryEvent> event) {
		if (!event.isOriginLocal()) {
			Map.Entry entry = event.getValue();
			if (entry != null) {
				this.executor.execute(() -> this.notifyListeners(event.getType(), entry));
			}
		}
		return CompletableFuture.completedStage(null);
	}

	@CacheEntryRemoved
	public CompletionStage removed(CacheEntryRemovedEvent> event) {
		if (!event.isOriginLocal()) {
			Map.Entry entry = event.getOldValue();
			// WFLY-4938 For some reason, the old value can be null
			if (entry != null) {
				this.executor.execute(() -> this.notifyListeners(event.getType(), entry));
			}
		}
		return CompletableFuture.completedStage(null);
	}

	private void notifyListeners(Event.Type type, Map.Entry entry) {
		this.notifyListeners(type, Collections.singletonMap(entry.getKey(), entry.getValue()));
	}

	private void notifyListeners(Event.Type type, Map entries) {
		for (Map.Entry, ExecutorService> entry: this.listeners.entrySet()) {
			RegistryListener listener = entry.getKey();
			Executor executor = entry.getValue();
			try {
				executor.execute(() -> {
					try {
						switch (type) {
							case CACHE_ENTRY_CREATED: {
								listener.added(entries);
								break;
							}
							case CACHE_ENTRY_MODIFIED: {
								listener.updated(entries);
								break;
							}
							case CACHE_ENTRY_REMOVED: {
								listener.removed(entries);
								break;
							}
							default: {
								throw new IllegalStateException(type.name());
							}
						}
					} catch (Throwable e) {
						LOGGER.warn(e.getLocalizedMessage(), e);
					}
				});
			} catch (RejectedExecutionException e) {
				// Executor was shutdown
			}
		}
	}

	private void shutdown(ExecutorService executor) {
		executor.shutdown();
		try {
			executor.awaitTermination(this.cache.getCacheConfiguration().transaction().cacheStopTimeout(), TimeUnit.MILLISECONDS);
		} catch (InterruptedException e) {
			Thread.currentThread().interrupt();
		}
	}
}