聊聊sharding-jdbc的XATransactionManager

public interface XATransactionManager extends AutoCloseable {
    
    /**
     * Initialize XA transaction manager.
     */
    void init();
    
    /**
     * Register recovery resource.
     *
     * @param dataSourceName data source name
     * @param xaDataSource XA data source
     */
    void registerRecoveryResource(String dataSourceName, XADataSource xaDataSource);
    
    /**
     * Remove recovery resource.
     *
     * @param dataSourceName data source name
     * @param xaDataSource   XA data source
     */
    void removeRecoveryResource(String dataSourceName, XADataSource xaDataSource);
    
    /**
     * Enlist single XA resource.
     * 
     * @param singleXAResource single XA resource
     */
    void enlistResource(SingleXAResource singleXAResource);
    
    /**
     * Get transaction manager.
     * 
     * @return transaction manager
     */
    TransactionManager getTransactionManager();
}

public final class AtomikosTransactionManager implements XATransactionManager {
    
    private final UserTransactionManager transactionManager = new UserTransactionManager();
    
    private final UserTransactionService userTransactionService = new UserTransactionServiceImp();
    
    @Override
    public void init() {
        userTransactionService.init();
    }
    
    @Override
    public void registerRecoveryResource(final String dataSourceName, final XADataSource xaDataSource) {
        userTransactionService.registerResource(new AtomikosXARecoverableResource(dataSourceName, xaDataSource));
    }
    
    @Override
    public void removeRecoveryResource(final String dataSourceName, final XADataSource xaDataSource) {
        userTransactionService.removeResource(new AtomikosXARecoverableResource(dataSourceName, xaDataSource));
    }
    
    @Override
    @SneakyThrows
    public void enlistResource(final SingleXAResource xaResource) {
        transactionManager.getTransaction().enlistResource(xaResource);
    }
    
    @Override
    public TransactionManager getTransactionManager() {
        return transactionManager;
    }
    
    @Override
    public void close() {
        userTransactionService.shutdown(true);
    }
}

@RunWith(MockitoJUnitRunner.class)
public final class AtomikosTransactionManagerTest {
    
    private AtomikosTransactionManager atomikosTransactionManager = new AtomikosTransactionManager();
    
    @Mock
    private UserTransactionManager userTransactionManager;
    
    @Mock
    private UserTransactionService userTransactionService;
    
    @Mock
    private XADataSource xaDataSource;
    
    @Before
    public void setUp() {
        ReflectiveUtil.setProperty(atomikosTransactionManager, "transactionManager", userTransactionManager);
        ReflectiveUtil.setProperty(atomikosTransactionManager, "userTransactionService", userTransactionService);
    }
    
    @Test
    public void assertInit() {
        atomikosTransactionManager.init();
        verify(userTransactionService).init();
    }
    
    @Test
    public void assertRegisterRecoveryResource() {
        atomikosTransactionManager.registerRecoveryResource("ds1", xaDataSource);
        verify(userTransactionService).registerResource(any(AtomikosXARecoverableResource.class));
    }
    
    @Test
    public void assertRemoveRecoveryResource() {
        atomikosTransactionManager.removeRecoveryResource("ds1", xaDataSource);
        verify(userTransactionService).removeResource(any(AtomikosXARecoverableResource.class));
    }
    
    @Test
    public void assertEnListResource() throws Exception {
        SingleXAResource singleXAResource = mock(SingleXAResource.class);
        Transaction transaction = mock(Transaction.class);
        when(userTransactionManager.getTransaction()).thenReturn(transaction);
        atomikosTransactionManager.enlistResource(singleXAResource);
        verify(transaction).enlistResource(singleXAResource);
    }
    
    @Test
    public void assertTransactionManager() {
        assertThat(atomikosTransactionManager.getTransactionManager(), CoreMatchers.<TransactionManager>is(userTransactionManager));
    }
    
    @Test
    public void assertClose() {
        atomikosTransactionManager.close();
        verify(userTransactionService).shutdown(true);
    }
}