solr/core/src/java/org/apache/solr/store/blockcache/BlockDirectory.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.io.File;
 import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.lang.invoke.MethodHandles;
 import java.util.Set;

 import org.apache.lucene.index.IndexFileNames;
 import org.apache.lucene.store.Directory;
 import org.apache.lucene.store.FSDirectory;
 import org.apache.lucene.store.FilterDirectory;
 import org.apache.lucene.store.IOContext;
 import org.apache.lucene.store.IndexInput;
 import org.apache.lucene.store.IndexOutput;
 import org.apache.solr.core.ShutdownAwareDirectory;
 import org.apache.solr.store.hdfs.HdfsDirectory;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 /**
  * @lucene.experimental
  */
 public class BlockDirectory extends FilterDirectory implements ShutdownAwareDirectory {
   private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());

   public static final long BLOCK_SHIFT = Integer.getInteger("solr.hdfs.blockcache.blockshift", 13);

   public static final int BLOCK_SIZE = 1 << BLOCK_SHIFT;
   public static final long BLOCK_MOD = BLOCK_SIZE - 1;

   public static long getBlock(long pos) {
     return pos >>> BLOCK_SHIFT;
   }

   public static long getPosition(long pos) {
     return pos & BLOCK_MOD;
   }

   public static long getRealPosition(long block, long positionInBlock) {
     return (block << BLOCK_SHIFT) + positionInBlock;
   }

   public static Cache NO_CACHE = new Cache() {

     @Override
     public void update(String name, long blockId, int blockOffset,
         byte[] buffer, int offset, int length) {}

     @Override
     public boolean fetch(String name, long blockId, int blockOffset, byte[] b,
         int off, int lengthToReadInBlock) {
       return false;
     }

     @Override
     public void delete(String name) {

     }

     @Override
     public long size() {
       return 0;
     }

     @Override
     public void renameCacheFile(String source, String dest) {}

     @Override
     public void releaseResources() {}
   };

   private final int blockSize;
   private final String dirName;
   private final Cache cache;
   private final Set<String> blockCacheFileTypes;
   private final boolean blockCacheReadEnabled;
   private final boolean blockCacheWriteEnabled;

   private boolean cacheMerges;
   private boolean cacheReadOnce;

   public BlockDirectory(String dirName, Directory directory, Cache cache,
       Set<String> blockCacheFileTypes, boolean blockCacheReadEnabled,
       boolean blockCacheWriteEnabled) throws IOException {
     this(dirName, directory, cache, blockCacheFileTypes, blockCacheReadEnabled, blockCacheWriteEnabled, true, true);
   }

   public BlockDirectory(String dirName, Directory directory, Cache cache,
       Set<String> blockCacheFileTypes, boolean blockCacheReadEnabled,
       boolean blockCacheWriteEnabled, boolean cacheMerges, boolean cacheReadOnce) throws IOException {
     super(directory);
     this.cacheMerges = cacheMerges;
     this.cacheReadOnce = cacheReadOnce;
     this.dirName = dirName;
     blockSize = BLOCK_SIZE;
     this.cache = cache;
     if (blockCacheFileTypes == null || blockCacheFileTypes.isEmpty()) {
       this.blockCacheFileTypes = null;
     } else {
       this.blockCacheFileTypes = blockCacheFileTypes;
     }
     this.blockCacheReadEnabled = blockCacheReadEnabled;
     if (!blockCacheReadEnabled) {
       log.info("Block cache on read is disabled");
     }
     this.blockCacheWriteEnabled = blockCacheWriteEnabled;
     if (!blockCacheWriteEnabled) {
       log.info("Block cache on write is disabled");
     }
   }

   private IndexInput openInput(String name, int bufferSize, IOContext context)
       throws IOException {
     final IndexInput source = super.openInput(name, context);
     if (useReadCache(name, context)) {
       return new CachedIndexInput(source, blockSize, name,
           getFileCacheName(name), cache, bufferSize);
     }
     return source;
   }

   private boolean isCachableFile(String name) {
     for (String ext : blockCacheFileTypes) {
       if (name.endsWith(ext)) {
         return true;
       }
     }
     return false;
   }

   @Override
   public IndexInput openInput(final String name, IOContext context)
       throws IOException {
     return openInput(name, blockSize, context);
   }

   static class CachedIndexInput extends CustomBufferedIndexInput {
     private final Store store;
     private IndexInput source;
     private final int blockSize;
     private final long fileLength;
     private final String cacheName;
     private final Cache cache;

     public CachedIndexInput(IndexInput source, int blockSize, String name,
         String cacheName, Cache cache, int bufferSize) {
       super(name, bufferSize);
       this.source = source;
       this.blockSize = blockSize;
       fileLength = source.length();
       this.cacheName = cacheName;
       this.cache = cache;
       store = BufferStore.instance(blockSize);
     }

     @Override
     public IndexInput clone() {
       CachedIndexInput clone = (CachedIndexInput) super.clone();
       clone.source = source.clone();
       return clone;
     }

     @Override
     public long length() {
       return source.length();
     }

     @Override
     protected void seekInternal(long pos) throws IOException {}

     @Override
     protected void readInternal(byte[] b, int off, int len) throws IOException {
       long position = getFilePointer();
       while (len > 0) {
         int length = fetchBlock(position, b, off, len);
         position += length;
         len -= length;
         off += length;
       }
     }

     private int fetchBlock(long position, byte[] b, int off, int len)
         throws IOException {
       // read whole block into cache and then provide needed data
       long blockId = getBlock(position);
       int blockOffset = (int) getPosition(position);
       int lengthToReadInBlock = Math.min(len, blockSize - blockOffset);
       if (checkCache(blockId, blockOffset, b, off, lengthToReadInBlock)) {
         return lengthToReadInBlock;
       } else {
         readIntoCacheAndResult(blockId, blockOffset, b, off,
             lengthToReadInBlock);
       }
       return lengthToReadInBlock;
     }

     private void readIntoCacheAndResult(long blockId, int blockOffset,
         byte[] b, int off, int lengthToReadInBlock) throws IOException {
       long position = getRealPosition(blockId, 0);
       int length = (int) Math.min(blockSize, fileLength - position);
       source.seek(position);

       byte[] buf = store.takeBuffer(blockSize);
       source.readBytes(buf, 0, length);
       System.arraycopy(buf, blockOffset, b, off, lengthToReadInBlock);
       cache.update(cacheName, blockId, 0, buf, 0, blockSize);
       store.putBuffer(buf);
     }

     private boolean checkCache(long blockId, int blockOffset, byte[] b,
         int off, int lengthToReadInBlock) {
       return cache.fetch(cacheName, blockId, blockOffset, b, off,
           lengthToReadInBlock);
     }

     @Override
     protected void closeInternal() throws IOException {
       source.close();
     }
   }

   @Override
   public void closeOnShutdown() throws IOException {
     log.info("BlockDirectory closing on shutdown");
     // we are shutting down, no need to clean up cache
     super.close();
   }

   @Override
   public void close() throws IOException {
     try {
       String[] files = listAll();

       for (String file : files) {
         cache.delete(getFileCacheName(file));
       }

     } catch (FileNotFoundException e) {
       // the local file system folder may be gone
     } finally {
       super.close();
       cache.releaseResources();
     }
   }

   String getFileCacheName(String name) throws IOException {
     return getFileCacheLocation(name) + ":" + getFileModified(name);
   }

   private long getFileModified(String name) throws IOException {
     if (in instanceof FSDirectory) {
       File directory = ((FSDirectory) in).getDirectory().toFile();
       File file = new File(directory, name);
       if (!file.exists()) {
         throw new FileNotFoundException("File [" + name + "] not found");
       }
       return file.lastModified();
     } else if (in instanceof HdfsDirectory) {
       return ((HdfsDirectory) in).fileModified(name);
     } else {
       throw new UnsupportedOperationException();
     }
   }

   String getFileCacheLocation(String name) {
     return dirName + "/" + name;
   }

   /**
    * Expert: mostly for tests
    *
    * @lucene.experimental
    */
   public Cache getCache() {
     return cache;
   }

   /**
    * Determine whether read caching should be used for a particular
    * file/context.
    */
   boolean useReadCache(String name, IOContext context) {
     if (!blockCacheReadEnabled) {
       return false;
     }
     if (blockCacheFileTypes != null && !isCachableFile(name)) {
       return false;
     }
     switch (context.context) {
       // depending on params, we don't cache on merges or when only reading once
       case MERGE: {
         return cacheMerges;
       }
       case READ: {
         if (context.readOnce) {
           return cacheReadOnce;
         } else {
           return true;
         }
       }
       default: {
         return true;
       }
     }
   }

   /**
    * Determine whether write caching should be used for a particular
    * file/context.
    */
   boolean useWriteCache(String name, IOContext context) {
     if (!blockCacheWriteEnabled || name.startsWith(IndexFileNames.PENDING_SEGMENTS)) {
       // for safety, don't bother caching pending commits.
       // the cache does support renaming (renameCacheFile), but thats a scary optimization.
       return false;
     }
     if (blockCacheFileTypes != null && !isCachableFile(name)) {
       return false;
     }
     switch (context.context) {
       case MERGE: {
         // we currently don't cache any merge context writes
         return false;
       }
       default: {
         return true;
       }
     }
   }

   @Override
   public IndexOutput createOutput(String name, IOContext context)
       throws IOException {
     final IndexOutput dest = super.createOutput(name, context);
     if (useWriteCache(name, context)) {
       return new CachedIndexOutput(this, dest, blockSize, name, cache, blockSize);
     }
     return dest;
   }

   public void deleteFile(String name) throws IOException {
     cache.delete(getFileCacheName(name));
     super.deleteFile(name);
   }

   public boolean isBlockCacheReadEnabled() {
     return blockCacheReadEnabled;
   }

   public boolean isBlockCacheWriteEnabled() {
     return blockCacheWriteEnabled;
   }

 }
	/*
	* 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.io.File;
	import java.io.FileNotFoundException;
	import java.io.IOException;
	import java.lang.invoke.MethodHandles;
	import java.util.Set;

	import org.apache.lucene.index.IndexFileNames;
	import org.apache.lucene.store.Directory;
	import org.apache.lucene.store.FSDirectory;
	import org.apache.lucene.store.FilterDirectory;
	import org.apache.lucene.store.IOContext;
	import org.apache.lucene.store.IndexInput;
	import org.apache.lucene.store.IndexOutput;
	import org.apache.solr.core.ShutdownAwareDirectory;
	import org.apache.solr.store.hdfs.HdfsDirectory;
	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	/**
	* @lucene.experimental
	*/
	public class BlockDirectory extends FilterDirectory implements ShutdownAwareDirectory {
	private static final Logger log = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());

	public static final long BLOCK_SHIFT = Integer.getInteger("solr.hdfs.blockcache.blockshift", 13);

	public static final int BLOCK_SIZE = 1 << BLOCK_SHIFT;
	public static final long BLOCK_MOD = BLOCK_SIZE - 1;

	public static long getBlock(long pos) {
	return pos >>> BLOCK_SHIFT;
	}

	public static long getPosition(long pos) {
	return pos & BLOCK_MOD;
	}

	public static long getRealPosition(long block, long positionInBlock) {
	return (block << BLOCK_SHIFT) + positionInBlock;
	}

	public static Cache NO_CACHE = new Cache() {

	@Override
	public void update(String name, long blockId, int blockOffset,
	byte[] buffer, int offset, int length) {}

	@Override
	public boolean fetch(String name, long blockId, int blockOffset, byte[] b,
	int off, int lengthToReadInBlock) {
	return false;
	}

	@Override
	public void delete(String name) {

	}

	@Override
	public long size() {
	return 0;
	}

	@Override
	public void renameCacheFile(String source, String dest) {}

	@Override
	public void releaseResources() {}
	};

	private final int blockSize;
	private final String dirName;
	private final Cache cache;
	private final Set<String> blockCacheFileTypes;
	private final boolean blockCacheReadEnabled;
	private final boolean blockCacheWriteEnabled;

	private boolean cacheMerges;
	private boolean cacheReadOnce;

	public BlockDirectory(String dirName, Directory directory, Cache cache,
	Set<String> blockCacheFileTypes, boolean blockCacheReadEnabled,
	boolean blockCacheWriteEnabled) throws IOException {
	this(dirName, directory, cache, blockCacheFileTypes, blockCacheReadEnabled, blockCacheWriteEnabled, true, true);
	}

	public BlockDirectory(String dirName, Directory directory, Cache cache,
	Set<String> blockCacheFileTypes, boolean blockCacheReadEnabled,
	boolean blockCacheWriteEnabled, boolean cacheMerges, boolean cacheReadOnce) throws IOException {
	super(directory);
	this.cacheMerges = cacheMerges;
	this.cacheReadOnce = cacheReadOnce;
	this.dirName = dirName;
	blockSize = BLOCK_SIZE;
	this.cache = cache;
	if (blockCacheFileTypes == null \|\| blockCacheFileTypes.isEmpty()) {
	this.blockCacheFileTypes = null;
	} else {
	this.blockCacheFileTypes = blockCacheFileTypes;
	}
	this.blockCacheReadEnabled = blockCacheReadEnabled;
	if (!blockCacheReadEnabled) {
	log.info("Block cache on read is disabled");
	}
	this.blockCacheWriteEnabled = blockCacheWriteEnabled;
	if (!blockCacheWriteEnabled) {
	log.info("Block cache on write is disabled");
	}
	}

	private IndexInput openInput(String name, int bufferSize, IOContext context)
	throws IOException {
	final IndexInput source = super.openInput(name, context);
	if (useReadCache(name, context)) {
	return new CachedIndexInput(source, blockSize, name,
	getFileCacheName(name), cache, bufferSize);
	}
	return source;
	}

	private boolean isCachableFile(String name) {
	for (String ext : blockCacheFileTypes) {
	if (name.endsWith(ext)) {
	return true;
	}
	}
	return false;
	}

	@Override
	public IndexInput openInput(final String name, IOContext context)
	throws IOException {
	return openInput(name, blockSize, context);
	}

	static class CachedIndexInput extends CustomBufferedIndexInput {
	private final Store store;
	private IndexInput source;
	private final int blockSize;
	private final long fileLength;
	private final String cacheName;
	private final Cache cache;

	public CachedIndexInput(IndexInput source, int blockSize, String name,
	String cacheName, Cache cache, int bufferSize) {
	super(name, bufferSize);
	this.source = source;
	this.blockSize = blockSize;
	fileLength = source.length();
	this.cacheName = cacheName;
	this.cache = cache;
	store = BufferStore.instance(blockSize);
	}

	@Override
	public IndexInput clone() {
	CachedIndexInput clone = (CachedIndexInput) super.clone();
	clone.source = source.clone();
	return clone;
	}

	@Override
	public long length() {
	return source.length();
	}

	@Override
	protected void seekInternal(long pos) throws IOException {}

	@Override
	protected void readInternal(byte[] b, int off, int len) throws IOException {
	long position = getFilePointer();
	while (len > 0) {
	int length = fetchBlock(position, b, off, len);
	position += length;
	len -= length;
	off += length;
	}
	}

	private int fetchBlock(long position, byte[] b, int off, int len)
	throws IOException {
	// read whole block into cache and then provide needed data
	long blockId = getBlock(position);
	int blockOffset = (int) getPosition(position);
	int lengthToReadInBlock = Math.min(len, blockSize - blockOffset);
	if (checkCache(blockId, blockOffset, b, off, lengthToReadInBlock)) {
	return lengthToReadInBlock;
	} else {
	readIntoCacheAndResult(blockId, blockOffset, b, off,
	lengthToReadInBlock);
	}
	return lengthToReadInBlock;
	}

	private void readIntoCacheAndResult(long blockId, int blockOffset,
	byte[] b, int off, int lengthToReadInBlock) throws IOException {
	long position = getRealPosition(blockId, 0);
	int length = (int) Math.min(blockSize, fileLength - position);
	source.seek(position);

	byte[] buf = store.takeBuffer(blockSize);
	source.readBytes(buf, 0, length);
	System.arraycopy(buf, blockOffset, b, off, lengthToReadInBlock);
	cache.update(cacheName, blockId, 0, buf, 0, blockSize);
	store.putBuffer(buf);
	}

	private boolean checkCache(long blockId, int blockOffset, byte[] b,
	int off, int lengthToReadInBlock) {
	return cache.fetch(cacheName, blockId, blockOffset, b, off,
	lengthToReadInBlock);
	}

	@Override
	protected void closeInternal() throws IOException {
	source.close();
	}
	}

	@Override
	public void closeOnShutdown() throws IOException {
	log.info("BlockDirectory closing on shutdown");
	// we are shutting down, no need to clean up cache
	super.close();
	}

	@Override
	public void close() throws IOException {
	try {
	String[] files = listAll();

	for (String file : files) {
	cache.delete(getFileCacheName(file));
	}

	} catch (FileNotFoundException e) {
	// the local file system folder may be gone
	} finally {
	super.close();
	cache.releaseResources();
	}
	}

	String getFileCacheName(String name) throws IOException {
	return getFileCacheLocation(name) + ":" + getFileModified(name);
	}

	private long getFileModified(String name) throws IOException {
	if (in instanceof FSDirectory) {
	File directory = ((FSDirectory) in).getDirectory().toFile();
	File file = new File(directory, name);
	if (!file.exists()) {
	throw new FileNotFoundException("File [" + name + "] not found");
	}
	return file.lastModified();
	} else if (in instanceof HdfsDirectory) {
	return ((HdfsDirectory) in).fileModified(name);
	} else {
	throw new UnsupportedOperationException();
	}
	}

	String getFileCacheLocation(String name) {
	return dirName + "/" + name;
	}

	/**
	* Expert: mostly for tests
	*
	* @lucene.experimental
	*/
	public Cache getCache() {
	return cache;
	}

	/**
	* Determine whether read caching should be used for a particular
	* file/context.
	*/
	boolean useReadCache(String name, IOContext context) {
	if (!blockCacheReadEnabled) {
	return false;
	}
	if (blockCacheFileTypes != null && !isCachableFile(name)) {
	return false;
	}
	switch (context.context) {
	// depending on params, we don't cache on merges or when only reading once
	case MERGE: {
	return cacheMerges;
	}
	case READ: {
	if (context.readOnce) {
	return cacheReadOnce;
	} else {
	return true;
	}
	}
	default: {
	return true;
	}
	}
	}

	/**
	* Determine whether write caching should be used for a particular
	* file/context.
	*/
	boolean useWriteCache(String name, IOContext context) {
	if (!blockCacheWriteEnabled \|\| name.startsWith(IndexFileNames.PENDING_SEGMENTS)) {
	// for safety, don't bother caching pending commits.
	// the cache does support renaming (renameCacheFile), but thats a scary optimization.
	return false;
	}
	if (blockCacheFileTypes != null && !isCachableFile(name)) {
	return false;
	}
	switch (context.context) {
	case MERGE: {
	// we currently don't cache any merge context writes
	return false;
	}
	default: {
	return true;
	}
	}
	}

	@Override
	public IndexOutput createOutput(String name, IOContext context)
	throws IOException {
	final IndexOutput dest = super.createOutput(name, context);
	if (useWriteCache(name, context)) {
	return new CachedIndexOutput(this, dest, blockSize, name, cache, blockSize);
	}
	return dest;
	}

	public void deleteFile(String name) throws IOException {
	cache.delete(getFileCacheName(name));
	super.deleteFile(name);
	}

	public boolean isBlockCacheReadEnabled() {
	return blockCacheReadEnabled;
	}

	public boolean isBlockCacheWriteEnabled() {
	return blockCacheWriteEnabled;
	}

	}