solr/core/src/java/org/apache/solr/store/blockcache/BlockCache.java - lucene-solr - Git at Google

 /*
  * 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.solr.store.blockcache;

 import java.nio.ByteBuffer;
 import java.util.List;
 import java.util.concurrent.CopyOnWriteArrayList;
 import java.util.concurrent.atomic.AtomicInteger;

 import com.github.benmanes.caffeine.cache.Cache;
 import com.github.benmanes.caffeine.cache.Caffeine;
 import com.github.benmanes.caffeine.cache.RemovalCause;
 import com.github.benmanes.caffeine.cache.RemovalListener;

 /**
  * @lucene.experimental
  */
 public class BlockCache {

   public static final int _128M = 134217728;
   public static final int _32K = 32768;
   private final Cache<BlockCacheKey,BlockCacheLocation> cache;
   private final ByteBuffer[] banks;
   private final BlockLocks[] locks;
   private final AtomicInteger[] lockCounters;
   private final int blockSize;
   private final int numberOfBlocksPerBank;
   private final int maxEntries;
   private final Metrics metrics;
   private final List<OnRelease> onReleases = new CopyOnWriteArrayList<>();

   public static interface OnRelease {
     public void release(BlockCacheKey blockCacheKey);
   }

   public BlockCache(Metrics metrics, boolean directAllocation, long totalMemory) {
     this(metrics, directAllocation, totalMemory, _128M);
   }

   public BlockCache(Metrics metrics, boolean directAllocation,
       long totalMemory, int slabSize) {
     this(metrics, directAllocation, totalMemory, slabSize, _32K);
   }

   public BlockCache(Metrics metrics, boolean directAllocation,
       long totalMemory, int slabSize, int blockSize) {
     this.metrics = metrics;
     numberOfBlocksPerBank = slabSize / blockSize;
     int numberOfBanks = (int) (totalMemory / slabSize);

     banks = new ByteBuffer[numberOfBanks];
     locks = new BlockLocks[numberOfBanks];
     lockCounters = new AtomicInteger[numberOfBanks];
     maxEntries = (numberOfBlocksPerBank * numberOfBanks) - 1;
     for (int i = 0; i < numberOfBanks; i++) {
       if (directAllocation) {
         banks[i] = ByteBuffer.allocateDirect(numberOfBlocksPerBank * blockSize);
       } else {
         banks[i] = ByteBuffer.allocate(numberOfBlocksPerBank * blockSize);
       }
       locks[i] = new BlockLocks(numberOfBlocksPerBank);
       lockCounters[i] = new AtomicInteger();
     }

     RemovalListener<BlockCacheKey,BlockCacheLocation> listener = (blockCacheKey, blockCacheLocation, removalCause) -> releaseLocation(blockCacheKey, blockCacheLocation, removalCause);

     cache = Caffeine.newBuilder()
         .removalListener(listener)
         .maximumSize(maxEntries)
         .build();
     this.blockSize = blockSize;
   }

   public void release(BlockCacheKey key) {
     cache.invalidate(key);
   }

   private void releaseLocation(BlockCacheKey blockCacheKey, BlockCacheLocation location, RemovalCause removalCause) {
     if (location == null) {
       return;
     }
     int bankId = location.getBankId();
     int block = location.getBlock();

     // mark the block removed before we release the lock to allow it to be reused
     location.setRemoved(true);

     locks[bankId].clear(block);
     lockCounters[bankId].decrementAndGet();
     for (OnRelease onRelease : onReleases) {
       onRelease.release(blockCacheKey);
     }
     if (removalCause.wasEvicted()) {
       metrics.blockCacheEviction.incrementAndGet();
     }
     metrics.blockCacheSize.decrementAndGet();
   }

   /**
    * This is only best-effort... it's possible for false to be returned, meaning the block was not able to be cached.
    * NOTE: blocks may not currently be updated (false will be returned if the block is already cached)
    * The blockCacheKey is cloned before it is inserted into the map, so it may be reused by clients if desired.
    *
    * @param blockCacheKey the key for the block
    * @param blockOffset the offset within the block
    * @param data source data to write to the block
    * @param offset offset within the source data array
    * @param length the number of bytes to write.
    * @return true if the block was cached/updated
    */
   public boolean store(BlockCacheKey blockCacheKey, int blockOffset,
       byte[] data, int offset, int length) {
     if (length + blockOffset > blockSize) {
       throw new RuntimeException("Buffer size exceeded, expecting max ["
           + blockSize + "] got length [" + length + "] with blockOffset ["
           + blockOffset + "]");
     }
     BlockCacheLocation location = cache.getIfPresent(blockCacheKey);
     if (location == null) {
       location = new BlockCacheLocation();
       if (!findEmptyLocation(location)) {
         // YCS: it looks like when the cache is full (a normal scenario), then two concurrent writes will result in one of them failing
         // because no eviction is done first.  The code seems to rely on leaving just a single block empty.
         // TODO: simplest fix would be to leave more than one block empty
         metrics.blockCacheStoreFail.incrementAndGet();
         return false;
       }
     } else {
       // If we allocated a new block, then it has never been published and is thus never in danger of being concurrently removed.
       // On the other hand, if this is an existing block we are updating, it may concurrently be removed and reused for another
       // purpose (and then our write may overwrite that).  This can happen even if clients never try to update existing blocks,
       // since two clients can try to cache the same block concurrently.  Because of this, the ability to update an existing
       // block has been removed for the time being (see SOLR-10121).

       // No metrics to update: we don't count a redundant store as a store fail.
       return false;
     }

     int bankId = location.getBankId();
     int bankOffset = location.getBlock() * blockSize;
     ByteBuffer bank = getBank(bankId);
     bank.position(bankOffset + blockOffset);
     bank.put(data, offset, length);

     // make sure all modifications to the block have been completed before we publish it.
     cache.put(blockCacheKey.clone(), location);
     metrics.blockCacheSize.incrementAndGet();
     return true;
   }

   /**
    * @param blockCacheKey the key for the block
    * @param buffer the target buffer for the read result
    * @param blockOffset offset within the block
    * @param off offset within the target buffer
    * @param length the number of bytes to read
    * @return true if the block was cached and the bytes were read
    */
   public boolean fetch(BlockCacheKey blockCacheKey, byte[] buffer,
       int blockOffset, int off, int length) {
     BlockCacheLocation location = cache.getIfPresent(blockCacheKey);
     if (location == null) {
       metrics.blockCacheMiss.incrementAndGet();
       return false;
     }

     int bankId = location.getBankId();
     int bankOffset = location.getBlock() * blockSize;
     location.touch();
     ByteBuffer bank = getBank(bankId);
     bank.position(bankOffset + blockOffset);
     bank.get(buffer, off, length);

     if (location.isRemoved()) {
       // must check *after* the read is done since the bank may have been reused for another block
       // before or during the read.
       metrics.blockCacheMiss.incrementAndGet();
       return false;
     }

     metrics.blockCacheHit.incrementAndGet();
     return true;
   }

   public boolean fetch(BlockCacheKey blockCacheKey, byte[] buffer) {
     checkLength(buffer);
     return fetch(blockCacheKey, buffer, 0, 0, blockSize);
   }

   private boolean findEmptyLocation(BlockCacheLocation location) {
     // This is a tight loop that will try and find a location to
     // place the block before giving up
     for (int j = 0; j < 10; j++) {
       OUTER: for (int bankId = 0; bankId < banks.length; bankId++) {
         AtomicInteger bitSetCounter = lockCounters[bankId];
         BlockLocks bitSet = locks[bankId];
         if (bitSetCounter.get() == numberOfBlocksPerBank) {
           // if bitset is full
           continue OUTER;
         }
         // this check needs to spin, if a lock was attempted but not obtained
         // the rest of the bank should not be skipped
         int bit = bitSet.nextClearBit(0);
         INNER: while (bit != -1) {
           if (bit >= numberOfBlocksPerBank) {
             // bit set is full
             continue OUTER;
           }
           if (!bitSet.set(bit)) {
             // lock was not obtained
             // this restarts at 0 because another block could have been unlocked
             // while this was executing
             bit = bitSet.nextClearBit(0);
             continue INNER;
           } else {
             // lock obtained
             location.setBankId(bankId);
             location.setBlock(bit);
             bitSetCounter.incrementAndGet();
             return true;
           }
         }
       }
     }
     return false;
   }

   private void checkLength(byte[] buffer) {
     if (buffer.length != blockSize) {
       throw new RuntimeException("Buffer wrong size, expecting [" + blockSize
           + "] got [" + buffer.length + "]");
     }
   }

   /** Returns a new copy of the ByteBuffer for the given bank, so it's safe to call position() on w/o additional synchronization */
   private ByteBuffer getBank(int bankId) {
     return banks[bankId].duplicate();
   }

   /** returns the number of elements in the cache */
   public int getSize() {
     return cache.asMap().size();
   }

   void setOnRelease(OnRelease onRelease) {
     this.onReleases.add(onRelease);
   }
 }
	/*
	* 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.solr.store.blockcache;

	import java.nio.ByteBuffer;
	import java.util.List;
	import java.util.concurrent.CopyOnWriteArrayList;
	import java.util.concurrent.atomic.AtomicInteger;

	import com.github.benmanes.caffeine.cache.Cache;
	import com.github.benmanes.caffeine.cache.Caffeine;
	import com.github.benmanes.caffeine.cache.RemovalCause;
	import com.github.benmanes.caffeine.cache.RemovalListener;

	/**
	* @lucene.experimental
	*/
	public class BlockCache {

	public static final int _128M = 134217728;
	public static final int _32K = 32768;
	private final Cache<BlockCacheKey,BlockCacheLocation> cache;
	private final ByteBuffer[] banks;
	private final BlockLocks[] locks;
	private final AtomicInteger[] lockCounters;
	private final int blockSize;
	private final int numberOfBlocksPerBank;
	private final int maxEntries;
	private final Metrics metrics;
	private final List<OnRelease> onReleases = new CopyOnWriteArrayList<>();

	public static interface OnRelease {
	public void release(BlockCacheKey blockCacheKey);
	}

	public BlockCache(Metrics metrics, boolean directAllocation, long totalMemory) {
	this(metrics, directAllocation, totalMemory, _128M);
	}

	public BlockCache(Metrics metrics, boolean directAllocation,
	long totalMemory, int slabSize) {
	this(metrics, directAllocation, totalMemory, slabSize, _32K);
	}

	public BlockCache(Metrics metrics, boolean directAllocation,
	long totalMemory, int slabSize, int blockSize) {
	this.metrics = metrics;
	numberOfBlocksPerBank = slabSize / blockSize;
	int numberOfBanks = (int) (totalMemory / slabSize);

	banks = new ByteBuffer[numberOfBanks];
	locks = new BlockLocks[numberOfBanks];
	lockCounters = new AtomicInteger[numberOfBanks];
	maxEntries = (numberOfBlocksPerBank * numberOfBanks) - 1;
	for (int i = 0; i < numberOfBanks; i++) {
	if (directAllocation) {
	banks[i] = ByteBuffer.allocateDirect(numberOfBlocksPerBank * blockSize);
	} else {
	banks[i] = ByteBuffer.allocate(numberOfBlocksPerBank * blockSize);
	}
	locks[i] = new BlockLocks(numberOfBlocksPerBank);
	lockCounters[i] = new AtomicInteger();
	}

	RemovalListener<BlockCacheKey,BlockCacheLocation> listener = (blockCacheKey, blockCacheLocation, removalCause) -> releaseLocation(blockCacheKey, blockCacheLocation, removalCause);

	cache = Caffeine.newBuilder()
	.removalListener(listener)
	.maximumSize(maxEntries)
	.build();
	this.blockSize = blockSize;
	}

	public void release(BlockCacheKey key) {
	cache.invalidate(key);
	}

	private void releaseLocation(BlockCacheKey blockCacheKey, BlockCacheLocation location, RemovalCause removalCause) {
	if (location == null) {
	return;
	}
	int bankId = location.getBankId();
	int block = location.getBlock();

	// mark the block removed before we release the lock to allow it to be reused
	location.setRemoved(true);

	locks[bankId].clear(block);
	lockCounters[bankId].decrementAndGet();
	for (OnRelease onRelease : onReleases) {
	onRelease.release(blockCacheKey);
	}
	if (removalCause.wasEvicted()) {
	metrics.blockCacheEviction.incrementAndGet();
	}
	metrics.blockCacheSize.decrementAndGet();
	}

	/**
	* This is only best-effort... it's possible for false to be returned, meaning the block was not able to be cached.
	* NOTE: blocks may not currently be updated (false will be returned if the block is already cached)
	* The blockCacheKey is cloned before it is inserted into the map, so it may be reused by clients if desired.
	*
	* @param blockCacheKey the key for the block
	* @param blockOffset the offset within the block
	* @param data source data to write to the block
	* @param offset offset within the source data array
	* @param length the number of bytes to write.
	* @return true if the block was cached/updated
	*/
	public boolean store(BlockCacheKey blockCacheKey, int blockOffset,
	byte[] data, int offset, int length) {
	if (length + blockOffset > blockSize) {
	throw new RuntimeException("Buffer size exceeded, expecting max ["
	+ blockSize + "] got length [" + length + "] with blockOffset ["
	+ blockOffset + "]");
	}
	BlockCacheLocation location = cache.getIfPresent(blockCacheKey);
	if (location == null) {
	location = new BlockCacheLocation();
	if (!findEmptyLocation(location)) {
	// YCS: it looks like when the cache is full (a normal scenario), then two concurrent writes will result in one of them failing
	// because no eviction is done first. The code seems to rely on leaving just a single block empty.
	// TODO: simplest fix would be to leave more than one block empty
	metrics.blockCacheStoreFail.incrementAndGet();
	return false;
	}
	} else {
	// If we allocated a new block, then it has never been published and is thus never in danger of being concurrently removed.
	// On the other hand, if this is an existing block we are updating, it may concurrently be removed and reused for another
	// purpose (and then our write may overwrite that). This can happen even if clients never try to update existing blocks,
	// since two clients can try to cache the same block concurrently. Because of this, the ability to update an existing
	// block has been removed for the time being (see SOLR-10121).

	// No metrics to update: we don't count a redundant store as a store fail.
	return false;
	}

	int bankId = location.getBankId();
	int bankOffset = location.getBlock() * blockSize;
	ByteBuffer bank = getBank(bankId);
	bank.position(bankOffset + blockOffset);
	bank.put(data, offset, length);

	// make sure all modifications to the block have been completed before we publish it.
	cache.put(blockCacheKey.clone(), location);
	metrics.blockCacheSize.incrementAndGet();
	return true;
	}

	/**
	* @param blockCacheKey the key for the block
	* @param buffer the target buffer for the read result
	* @param blockOffset offset within the block
	* @param off offset within the target buffer
	* @param length the number of bytes to read
	* @return true if the block was cached and the bytes were read
	*/
	public boolean fetch(BlockCacheKey blockCacheKey, byte[] buffer,
	int blockOffset, int off, int length) {
	BlockCacheLocation location = cache.getIfPresent(blockCacheKey);
	if (location == null) {
	metrics.blockCacheMiss.incrementAndGet();
	return false;
	}

	int bankId = location.getBankId();
	int bankOffset = location.getBlock() * blockSize;
	location.touch();
	ByteBuffer bank = getBank(bankId);
	bank.position(bankOffset + blockOffset);
	bank.get(buffer, off, length);

	if (location.isRemoved()) {
	// must check after the read is done since the bank may have been reused for another block
	// before or during the read.
	metrics.blockCacheMiss.incrementAndGet();
	return false;
	}

	metrics.blockCacheHit.incrementAndGet();
	return true;
	}

	public boolean fetch(BlockCacheKey blockCacheKey, byte[] buffer) {
	checkLength(buffer);
	return fetch(blockCacheKey, buffer, 0, 0, blockSize);
	}

	private boolean findEmptyLocation(BlockCacheLocation location) {
	// This is a tight loop that will try and find a location to
	// place the block before giving up
	for (int j = 0; j < 10; j++) {
	OUTER: for (int bankId = 0; bankId < banks.length; bankId++) {
	AtomicInteger bitSetCounter = lockCounters[bankId];
	BlockLocks bitSet = locks[bankId];
	if (bitSetCounter.get() == numberOfBlocksPerBank) {
	// if bitset is full
	continue OUTER;
	}
	// this check needs to spin, if a lock was attempted but not obtained
	// the rest of the bank should not be skipped
	int bit = bitSet.nextClearBit(0);
	INNER: while (bit != -1) {
	if (bit >= numberOfBlocksPerBank) {
	// bit set is full
	continue OUTER;
	}
	if (!bitSet.set(bit)) {
	// lock was not obtained
	// this restarts at 0 because another block could have been unlocked
	// while this was executing
	bit = bitSet.nextClearBit(0);
	continue INNER;
	} else {
	// lock obtained
	location.setBankId(bankId);
	location.setBlock(bit);
	bitSetCounter.incrementAndGet();
	return true;
	}
	}
	}
	}
	return false;
	}

	private void checkLength(byte[] buffer) {
	if (buffer.length != blockSize) {
	throw new RuntimeException("Buffer wrong size, expecting [" + blockSize
	+ "] got [" + buffer.length + "]");
	}
	}

	/** Returns a new copy of the ByteBuffer for the given bank, so it's safe to call position() on w/o additional synchronization */
	private ByteBuffer getBank(int bankId) {
	return banks[bankId].duplicate();
	}

	/** returns the number of elements in the cache */
	public int getSize() {
	return cache.asMap().size();
	}

	void setOnRelease(OnRelease onRelease) {
	this.onReleases.add(onRelease);
	}
	}