package com.ldbc.driver;
import com.google.common.collect.Ordering;
import com.ldbc.driver.control.LoggingService;
import com.ldbc.driver.util.ClassLoaderHelper;
import java.io.Closeable;
import java.io.IOException;
import java.util.HashMap;
import java.util.Map;
import java.util.concurrent.atomic.AtomicBoolean;
import static java.lang.String.format;
public abstract class Db implements Closeable
{
private boolean isInitialized = false;
private AtomicBoolean isShutdown = new AtomicBoolean( false );
private DbConnectionState dbConnectionState = null;
private Map<Class<? extends Operation>,OperationHandler> operationHandlers = new HashMap<>();
private OperationHandler[] operationHandlersArray = null;
private OperationHandlerRunnerFactory operationHandlerRunnableContextFactory = null;
synchronized public final void init(
Map<String,String> params,
LoggingService loggingService,
Map<Integer,Class<? extends Operation>> operationTypeToClassMapping )
throws DbException
{
if ( true == isInitialized )
{
throw new DbException( "DB may be initialized only once" );
}
onInit( params, loggingService );
dbConnectionState = getConnectionState();
operationHandlerRunnableContextFactory = new PoolingOperationHandlerRunnerFactory(
new InstantiatingOperationHandlerRunnerFactory()
);
operationHandlersArray = toOperationHandlerArray( operationTypeToClassMapping, operationHandlers );
operationHandlers = null;
isInitialized = true;
}
/**
* Called once to initialize state for DB client
*/
protected abstract void onInit( Map<String,String> properties, LoggingService loggingService ) throws DbException;
@Override
synchronized public final void close() throws IOException
{
if ( isShutdown.get() )
{
throw new IOException( "DB may be cleaned up only once" );
}
isShutdown.set( true );
onClose();
try
{
operationHandlerRunnableContextFactory.shutdown();
}
catch ( OperationException e )
{
throw new IOException( "Error shutting down operation handler runnable factory", e );
}
}
// TODO this is a temporary hack to support warmup more easily, because the runnable contexts need to be cleared
// TODO ultimately this would be done in another way
synchronized public final void reInit() throws DbException
{
try
{
operationHandlerRunnableContextFactory.shutdown();
}
catch ( OperationException e )
{
throw new DbException( "Error shutting down operation handler runnable factory", e );
}
operationHandlerRunnableContextFactory = new PoolingOperationHandlerRunnerFactory(
new InstantiatingOperationHandlerRunnerFactory()
);
}
/**
* Called once to cleanup state for DB client
*/
protected abstract void onClose() throws IOException;
public final <A extends Operation, H extends OperationHandler<A,?>> void registerOperationHandler(
Class<A> operationType, Class<H> operationHandlerType ) throws DbException
{
if ( operationHandlers.containsKey( operationType ) )
{
throw new DbException( format( "Client already has handler registered for %s", operationType.getClass() ) );
}
try
{
OperationHandler operationHandler = ClassLoaderHelper.loadOperationHandler( operationHandlerType );
operationHandlers.put( operationType, operationHandler );
}
catch ( OperationException e )
{
throw new DbException(
format( "%s could not instantiate instance of %s",
getClass().getSimpleName(),
operationHandlerType.getSimpleName()
),
e );
}
}
public final OperationHandlerRunnableContext getOperationHandlerRunnableContext( Operation operation )
throws DbException
{
OperationHandler operationHandler = operationHandlersArray[operation.type()];
if ( null == operationHandler )
{
throw new DbException( format( "No handler registered for %s", operation.getClass() ) );
}
try
{
OperationHandlerRunnableContext operationHandlerRunnableContext =
operationHandlerRunnableContextFactory.newOperationHandlerRunner();
operationHandlerRunnableContext.setOperationHandler( operationHandler );
operationHandlerRunnableContext.setDbConnectionState( dbConnectionState );
return operationHandlerRunnableContext;
}
catch ( Exception e )
{
throw new DbException( format( "Unable to instantiate handler for operation:\n%s", operation ), e );
}
}
private static OperationHandler[] toOperationHandlerArray(
Map<Integer,Class<? extends Operation>> operationTypeToClassMapping,
Map<Class<? extends Operation>,OperationHandler> operationHandlers ) throws DbException
{
if ( operationTypeToClassMapping.isEmpty() )
{
return new OperationHandler[]{};
}
else
{
int minOperationType = Ordering.<Integer>natural().min( operationTypeToClassMapping.keySet() );
if ( minOperationType < 0 )
{
throw new DbException( format( "Operation type code lower than 0: %s", minOperationType ) );
}
int maxOperationType = Ordering.<Integer>natural().max( operationTypeToClassMapping.keySet() );
OperationHandler[] operationHandlersArray = new OperationHandler[maxOperationType + 1];
for ( int i = 0; i < operationHandlersArray.length; i++ )
{
if ( operationTypeToClassMapping.containsKey( i ) )
{
operationHandlersArray[i] = operationHandlers.get( operationTypeToClassMapping.get( i ) );
}
}
return operationHandlersArray;
}
}
/**
* Should return any state related to the database connection that can be
* reused by all operation handlers
*/
protected abstract DbConnectionState getConnectionState() throws DbException;
}