package com.atlassian.messagequeue.internal.inmemory;
   
   import com.atlassian.messagequeue.Message;
   import com.atlassian.messagequeue.MessageRunnerNotRegisteredException;
   import com.atlassian.messagequeue.MessageRunnerKey;
   import com.atlassian.messagequeue.MessageRunnerService;
   import com.atlassian.messagequeue.internal.core.MessageRunnerRegistryHelper;
   import com.atlassian.messagequeue.registry.MessageRunnerRegistryService;
   import com.atlassian.messagequeue.registry.MessageRunner;
  import com.atlassian.workcontext.api.WorkContextDoorway;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  import java.util.Optional;
  import java.util.concurrent.ConcurrentHashMap;
  import java.util.concurrent.ExecutorService;
  
  import static java.util.Objects.requireNonNull;
  
  /**
   * Implementation of both {@link MessageRunnerService} and {@link MessageRunnerRegistryService}. Messages are stored in-memory.
   */
  public class InMemoryMessageRunnerService implements MessageRunnerService, MessageRunnerRegistryService, MessageRunnerRegistryHelper {
  
      private static final Logger log = LoggerFactory.getLogger(InMemoryMessageRunnerService.class);
  
      private final ConcurrentHashMap<MessageRunnerKey, MessageRunner> runners = new ConcurrentHashMap<>();
  
      private final ExecutorService executorService;
  
      public InMemoryMessageRunnerService(ExecutorService executorService) {
          this.executorService = requireNonNull(executorService);
      }
  
      @Override
      public void registerMessageRunner(MessageRunnerKey messageRunnerKey, MessageRunner messageRunner) {
          runners.put(messageRunnerKey, messageRunner);
      }
  
      @Override
      public void unregisterMessageRunner(MessageRunnerKey messageRunnerKey) {
          final MessageRunner prev = runners.remove(messageRunnerKey);
          if (null == prev) {
              log.debug("Attempting to remove unregistered runner {}", messageRunnerKey);
          }
      }
  
      @Override
      public void addMessage(Message message) {
          final Optional<String> payload = message.getPayload();
          final MessageRunnerKey runnerKey = message.getRunnerKey();
          final MessageRunner messageRunner = runners.get(runnerKey);
          if (null == messageRunner) {
              throw new MessageRunnerNotRegisteredException(runnerKey);
          }
          executorService.submit(() -> {
              try (WorkContextDoorway ignored = new WorkContextDoorway().open()) {
                  messageRunner.processMessage(new InMemoryMessageContext(payload.orElse(null)));
              } catch (Exception e) {
                  log.error("failed to process message for key: " + runnerKey, e);
              }
          });
      }
  
      @Override
      public Optional<MessageRunner> getMessageRunner(MessageRunnerKey messageRunnerKey) {
          return Optional.ofNullable(runners.get(messageRunnerKey));
      }
  }