HSEARCH-2014 Expose size of the indexes among statistics

engine/src/main/java/org/hibernate/search/jmx/impl/JMXRegistrar.java

@@ -207,6 +207,16 @@ public int getNumberOfIndexedEntities(String entity) {
 		public Map<String, Integer> indexedEntitiesCount() {
 			return delegate.indexedEntitiesCount();
 		}
+
+		@Override
+		public long getIndexSize(String indexName) {
+			return delegate.getIndexSize( indexName );
+		}
+
+		@Override
+		public Map<String, Long> indexSizes() {
+			return delegate.indexSizes();
+		}
 	}
 
 	/**

engine/src/main/java/org/hibernate/search/stat/Statistics.java

@@ -114,4 +114,21 @@ public interface Statistics {
 	 *         the map value is the document count.
 	 */
 	Map<String, Integer> indexedEntitiesCount();
+
+	/**
+	 * Returns the size of the index with the given name, in bytes.
+	 *
+	 * @param indexName the index name (which may differ from the entity name,
+	 *         depending on configuration settings)
+	 * @return the size of the given index, in bytes
+	 * @throws java.lang.IllegalArgumentException in case the index name is not valid
+	 */
+	long getIndexSize(String indexName);
+
+	/**
+	 * @return a map of all index names and the size of the corresponding index, in bytes.
+	 *         The map key is the index name (which may differ from the entity name, depending
+	 *         on configuration settings).
+	 */
+	Map<String, Long> indexSizes();
 }

engine/src/main/java/org/hibernate/search/stat/impl/StatisticsImpl.java

@@ -6,6 +6,7 @@
  */
 package org.hibernate.search.stat.impl;
 
+import java.io.FileNotFoundException;
 import java.io.IOException;
 import java.util.HashMap;
 import java.util.HashSet;
@@ -15,6 +16,7 @@
 import java.util.concurrent.locks.Lock;
 import java.util.concurrent.locks.ReadWriteLock;
 import java.util.concurrent.locks.ReentrantReadWriteLock;
+import java.util.stream.Collectors;
 
 import org.apache.lucene.index.IndexReader;
 import org.apache.lucene.index.Term;
@@ -24,9 +26,11 @@
 import org.apache.lucene.search.MatchAllDocsQuery;
 import org.apache.lucene.search.TermQuery;
 import org.apache.lucene.search.TopDocs;
-
+import org.apache.lucene.store.Directory;
 import org.hibernate.search.engine.ProjectionConstants;
 import org.hibernate.search.exception.SearchException;
+import org.hibernate.search.indexes.spi.DirectoryBasedIndexManager;
+import org.hibernate.search.indexes.spi.IndexManager;
 import org.hibernate.search.engine.Version;
 import org.hibernate.search.engine.integration.impl.ExtendedSearchIntegrator;
 import org.hibernate.search.engine.service.classloading.spi.ClassLoadingException;
@@ -247,6 +251,49 @@ private Class<?> getEntityClass(String entity) {
 		}
 		return clazz;
 	}
+
+	@Override
+	public long getIndexSize(String indexName) {
+		IndexManager indexManager = extendedIntegrator.getIndexManager( indexName );
+		if ( indexManager == null ) {
+			throw new IllegalArgumentException( "'" + indexName + "' is not a known index" );
+		}
+		return getIndexSize( indexManager );
+	}
+
+	@Override
+	public Map<String, Long> indexSizes() {
+		return extendedIntegrator.getIndexManagerHolder().getIndexManagers().stream()
+				.collect( Collectors.toMap( IndexManager::getIndexName, this::getIndexSize ) );
+	}
+
+	private long getIndexSize(IndexManager indexManager) {
+		if ( !( indexManager instanceof DirectoryBasedIndexManager ) ) {
+			throw new IllegalArgumentException( "Index '" + indexManager.getIndexName()
+					+ "' is not a Lucene index" );
+		}
+
+		DirectoryBasedIndexManager directoryBasedIndexManager = (DirectoryBasedIndexManager) indexManager;
+		Directory directory = directoryBasedIndexManager.getDirectoryProvider().getDirectory();
+
+		long totalSize = 0l;
+		try {
+			for ( String fileName : directory.listAll() ) {
+				try {
+					totalSize += directory.fileLength( fileName );
+				}
+				catch (FileNotFoundException ignored) {
+					// Ignore: the file was probably removed since the call to listAll
+				}
+			}
+		}
+		catch (IOException e) {
+			throw new SearchException( "Unexpected exception while computing size of index '"
+					+ indexManager.getIndexName() + "'", e );
+		}
+
+		return totalSize;
+	}
 }
 
 

orm/src/test/java/org/hibernate/search/test/statistics/StatisticsTest.java

@@ -38,14 +38,32 @@
 import org.junit.After;
 import org.junit.Test;
 import org.junit.experimental.categories.Category;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+import org.junit.runners.Parameterized.Parameters;
 
 
 /**
  * @author Yoann Rodiere
  */
 @Category(ElasticsearchSupportInProgress.class) // HSEARCH-2421 Support statistics with Elasticsearch
+@RunWith( Parameterized.class )
 public class StatisticsTest extends SearchTestBase {
 
+	@Parameters(name = "Directory provider {0}")
+	public static Object[] data() {
+		return new Object[] {
+			"ram",
+			"filesystem" // Mainly to test getIndexSize()
+		};
+	}
+
+	private final String directoryProviderName;
+
+	public StatisticsTest(String directoryProviderName) {
+		this.directoryProviderName = directoryProviderName;
+	}
+
 	@Override
 	public Class<?>[] getAnnotatedClasses() {
 		return new Class<?>[] { A.class, B.class };
@@ -54,6 +72,7 @@ public Class<?>[] getAnnotatedClasses() {
 	@Override
 	public void configure(java.util.Map<String,Object> settings) {
 		settings.put( Environment.GENERATE_STATS, Boolean.TRUE.toString() );
+		settings.put( "hibernate.search.default.directory_provider", directoryProviderName );
 	}
 
 	@After
@@ -250,6 +269,59 @@ public void objectLoading_singleClassQueryLoader_criteria_iterate() {
 		}
 	}
 
+	@Test
+	@TestForIssue(jiraKey = "HSEARCH-2014")
+	public void indexSize() {
+		long currentSizeForA = getStatistics().getIndexSize( A.INDEX_NAME );
+		long currentSizeForB = getStatistics().getIndexSize( B.INDEX_NAME );
+		// Don't assume anything about the size of an empty index; it may not be 0
+
+		Map<String, Long> indexSizes = getStatistics().indexSizes();
+		assertEquals( 2, indexSizes.size() );
+		assertEquals( (Long) currentSizeForA, indexSizes.get( A.INDEX_NAME ) );
+		assertEquals( (Long) currentSizeForB, indexSizes.get( B.INDEX_NAME ) );
+
+		try (Session s = openSession()) {
+			A entity = new A();
+			entity.id = 1L;
+			Transaction tx = s.beginTransaction();
+			s.persist( entity );
+			tx.commit();
+		}
+
+		long previousSizeForA = currentSizeForA;
+		long previousSizeForB = currentSizeForB;
+		currentSizeForA = getStatistics().getIndexSize( A.INDEX_NAME );
+		currentSizeForB = getStatistics().getIndexSize( B.INDEX_NAME );
+		assertTrue( currentSizeForA > previousSizeForA );
+		assertEquals( previousSizeForB, currentSizeForB );
+
+		indexSizes = getStatistics().indexSizes();
+		assertEquals( 2, indexSizes.size() );
+		assertEquals( (Long) currentSizeForA, indexSizes.get( A.INDEX_NAME ) );
+		assertEquals( (Long) currentSizeForB, indexSizes.get( B.INDEX_NAME ) );
+
+		try (Session s = openSession()) {
+			A entity = new A();
+			entity.id = 2L;
+			Transaction tx = s.beginTransaction();
+			s.persist( entity );
+			tx.commit();
+		}
+
+		previousSizeForA = currentSizeForA;
+		previousSizeForB = currentSizeForB;
+		currentSizeForA = getStatistics().getIndexSize( A.INDEX_NAME );
+		currentSizeForB = getStatistics().getIndexSize( B.INDEX_NAME );
+		assertTrue( currentSizeForA > previousSizeForA );
+		assertEquals( previousSizeForB, currentSizeForB );
+
+		indexSizes = getStatistics().indexSizes();
+		assertEquals( 2, indexSizes.size() );
+		assertEquals( (Long) currentSizeForA, indexSizes.get( A.INDEX_NAME ) );
+		assertEquals( (Long) currentSizeForB, indexSizes.get( B.INDEX_NAME ) );
+	}
+
 	private Statistics getStatistics() {
 		return getSearchFactory().getStatistics();
 	}
@@ -261,6 +333,8 @@ private Query matchAll() {
 	@Entity
 	@Indexed
 	private static class A {
+		public static final String INDEX_NAME = A.class.getName();
+
 		@Id
 		private Long id;
 
@@ -269,8 +343,10 @@ private static class A {
 	}
 
 	@Entity
-	@Indexed
+	@Indexed(index = B.INDEX_NAME)
 	private static class B {
+		public static final String INDEX_NAME = "B";
+
 		@Id
 		private Long id;