All Downloads are FREE. Search and download functionalities are using the official Maven repository.

org.infinispan.server.resp.Resp3Handler Maven / Gradle / Ivy

There is a newer version: 15.1.0.Dev03
Show newest version
package org.infinispan.server.resp;

import static org.infinispan.server.resp.RespConstants.CRLF;
import static org.infinispan.server.resp.RespConstants.CRLF_STRING;
import static org.infinispan.server.resp.RespConstants.NULL;

import java.util.Collection;
import java.util.List;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.ScheduledExecutorService;
import java.util.function.Consumer;
import java.util.stream.Collectors;

import org.infinispan.AdvancedCache;
import org.infinispan.Cache;
import org.infinispan.commons.dataconversion.MediaType;
import org.infinispan.context.Flag;
import org.infinispan.factories.GlobalComponentRegistry;
import org.infinispan.factories.KnownComponentNames;
import org.infinispan.multimap.impl.EmbeddedMultimapListCache;
import org.infinispan.multimap.impl.EmbeddedMultimapPairCache;
import org.infinispan.multimap.impl.EmbeddedMultimapSortedSetCache;
import org.infinispan.multimap.impl.EmbeddedSetCache;
import org.infinispan.security.AuthorizationManager;
import org.infinispan.security.AuthorizationPermission;
import org.infinispan.security.actions.SecurityActions;
import org.infinispan.server.resp.commands.Resp3Command;
import org.infinispan.util.concurrent.BlockingManager;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.ByteBufUtil;
import io.netty.channel.ChannelHandlerContext;

public class Resp3Handler extends Resp3AuthHandler {
   private static final byte[] CRLF_BYTES = CRLF_STRING.getBytes();
   protected AdvancedCache ignorePreviousValueCache;
   protected EmbeddedMultimapListCache listMultimap;
   protected EmbeddedMultimapPairCache mapMultimap;
   protected EmbeddedSetCache embeddedSetCache;
   protected EmbeddedMultimapSortedSetCache sortedSetMultimap;
   protected final ScheduledExecutorService scheduler;
   protected final BlockingManager blockingManager;

   private final MediaType valueMediaType;

   Resp3Handler(RespServer respServer, MediaType valueMediaType) {
      super(respServer);
      this.valueMediaType = valueMediaType;

      GlobalComponentRegistry gcr = SecurityActions.getGlobalComponentRegistry(cache.getCacheManager());
      this.scheduler = gcr.getComponent(ScheduledExecutorService.class, KnownComponentNames.TIMEOUT_SCHEDULE_EXECUTOR);
      this.blockingManager = gcr.getComponent(BlockingManager.class);
   }

   @Override
   public void setCache(AdvancedCache cache) {
      super.setCache(cache);
      ignorePreviousValueCache = cache.withFlags(Flag.SKIP_CACHE_LOAD, Flag.IGNORE_RETURN_VALUES);
      Cache toMultimap = cache.withMediaType(MediaType.APPLICATION_OCTET_STREAM, valueMediaType);
      listMultimap = new EmbeddedMultimapListCache<>(toMultimap);
      mapMultimap = new EmbeddedMultimapPairCache<>(toMultimap);
      embeddedSetCache = new EmbeddedSetCache<>(toMultimap);
      sortedSetMultimap = new EmbeddedMultimapSortedSetCache<>(toMultimap);
   }

   public EmbeddedMultimapListCache getListMultimap() {
      return listMultimap;
   }

   public EmbeddedMultimapPairCache getHashMapMultimap() {
      return mapMultimap;
   }

   public EmbeddedSetCache getEmbeddedSetCache() {
      return embeddedSetCache;
   }

   public EmbeddedMultimapSortedSetCache getSortedSeMultimap() {
      return sortedSetMultimap;
   }

   public ScheduledExecutorService getScheduler() {
      return scheduler;
   }

   public BlockingManager getBlockingManager() {
      return blockingManager;
   }

   @Override
   protected CompletionStage actualHandleRequest(ChannelHandlerContext ctx, RespCommand type,
         List arguments) {
      if (type instanceof Resp3Command) {
         Resp3Command resp3Command = (Resp3Command) type;
         return resp3Command.perform(this, ctx, arguments);
      }
      return super.actualHandleRequest(ctx, type, arguments);
   }

   protected static void handleLongResult(Long result, ByteBufPool alloc) {
      ByteBufferUtils.writeLong(result, alloc);
   }

   protected static void handleDoubleResult(Double result, ByteBufPool alloc) {
      // TODO: this can be optimized to avoid the String allocation
      if (result == null) {
         handleNullResult(alloc);
      } else {
         handleBulkAsciiResult(Double.toString(result), alloc);
      }
   }

   protected static void handleCollectionDoubleResult(Collection collection, ByteBufPool alloc) {
      if (collection == null) {
         handleNullResult(alloc);
      } else {
         writeArrayPrefix(collection.size(), alloc);
         for(Double d: collection) {
            if (d == null) {
               handleNullResult(alloc);
            } else{
               handleDoubleResult(d, alloc);
            }
         }
      }
   }

   protected static void handleCollectionLongResult(Collection collection, ByteBufPool alloc) {
      if (collection == null) {
         handleNullResult(alloc);
      } else {
         String result = "*" + collection.size() + CRLF_STRING
               + collection.stream().map(value -> ":" + value + CRLF_STRING).collect(Collectors.joining());
         ByteBufferUtils.stringToByteBufAscii(result, alloc);
      }
   }

   public static void handleBulkResult(CharSequence result, ByteBufPool alloc) {
      if (result == null) {
         handleNullResult(alloc);
      } else {
         int resultLength = ByteBufUtil.utf8Bytes(result);
         int resultSizeLength = ByteBufferUtils.stringSize(resultLength);
         ByteBuf buf = alloc.acquire(1 + resultSizeLength + 2 + resultLength + 2);
         buf.writeByte('$');
         ByteBufferUtils.setIntChars(resultLength, resultSizeLength, buf);
         buf.writeBytes(CRLF);
         ByteBufUtil.writeUtf8(buf, result);
         buf.writeBytes(CRLF);
      }
   }

   public static void handleBulkAsciiResult(CharSequence result, ByteBufPool alloc) {
      if (result == null) {
         handleNullResult(alloc);
      } else {
         int resultLength = result.length();
         int resultSizeLength = ByteBufferUtils.stringSize(resultLength);
         ByteBuf buf = alloc.acquire(1 + resultSizeLength + 2 + resultLength + 2);
         buf.writeByte('$');
         ByteBufferUtils.setIntChars(resultLength, resultSizeLength, buf);
         buf.writeBytes(CRLF);
         ByteBufUtil.writeAscii(buf, result);
         buf.writeBytes(CRLF);
      }
   }

   public static void handleCollectionBulkResult(Collection collection, ByteBufPool alloc) {
      if (collection == null) {
         handleNullResult(alloc);
         return;
      }
      int dataLength = collection.stream().mapToInt(ba -> ba.length + 5 + lenghtInChars(ba.length)).sum();
      var buffer = allocAndWriteLengthPrefix('*', collection.size(), alloc, dataLength);
      collection.forEach(wba -> writeBulkResult(wba, buffer));
   }

   private static void handleNullResult(ByteBufPool alloc) {
      alloc.acquire(NULL.length).writeBytes(NULL);
   }

   protected static void handleBulkResult(byte[] result, ByteBufPool alloc) {
      if (result == null) {
         handleNullResult(alloc);
         return;
      }
      var buffer = allocAndWriteLengthPrefix('$', result.length, alloc, result.length + 2);
      buffer.writeBytes(result);
      buffer.writeBytes(CRLF_BYTES);
   }

   private static void writeBulkResult(byte[] result, ByteBuf buffer) {
      writeLengthPrefix('$', result.length, buffer);
      buffer.writeBytes(result);
      buffer.writeBytes(CRLF_BYTES);
   }

   protected static void handleThrowable(ByteBufPool alloc, Throwable t) {
      Consumer writer = RespErrorUtil.handleException(t);
      if (writer != null) {
         writer.accept(alloc);
      } else {
         ByteBufferUtils.stringToByteBuf("-ERR " + extractExceptionMessage(t) + CRLF_STRING, alloc);
      }
   }

   private static String extractExceptionMessage(Throwable t) {
      Throwable r = t;
      while (r.getCause() != null) {
         r = r.getCause();
      }

      return r.getMessage();
   }

   public AdvancedCache ignorePreviousValuesCache() {
      return ignorePreviousValueCache;
   }

   public CompletionStage delegate(ChannelHandlerContext ctx,
         RespCommand command,
         List arguments) {
      return super.actualHandleRequest(ctx, command, arguments);
   }

   public void checkPermission(AuthorizationPermission authorizationPermission) {
      AuthorizationManager authorizationManager = cache.getAuthorizationManager();
      if (authorizationManager != null) {
         authorizationManager.checkPermission(authorizationPermission);
      }
   }

   public static void writeArrayPrefix(int size, ByteBufPool alloc) {
      allocAndWriteLengthPrefix('*', size, alloc, 0);
   }

   public static void writeMapPrefix(int size, ByteBufPool alloc) {
      allocAndWriteLengthPrefix('%', size, alloc, 0);
   }

   private static ByteBuf allocAndWriteLengthPrefix(char type, int size, ByteBufPool alloc, int additionalBytes) {
      int strLength = lenghtInChars(size);
      ByteBuf buffer = alloc.acquire(strLength + additionalBytes + 3);
      buffer.writeByte(type);
      ByteBufferUtils.setIntChars(size, strLength, buffer);
      buffer.writeBytes(CRLF_BYTES);
      return buffer;
   }

   private static int lenghtInChars(int size) {
      int strLength = size == 0 ? 1 : (int) Math.log10(size) + 1;
      return strLength;
   }

   private static ByteBuf writeLengthPrefix(char type, int size, ByteBuf buffer) {
      int strLength = lenghtInChars(size);
      buffer.writeByte(type);
      ByteBufferUtils.setIntChars(size, strLength, buffer);
      buffer.writeBytes(CRLF_BYTES);
      return buffer;
   }

}




© 2015 - 2024 Weber Informatics LLC | Privacy Policy