> mvccEntryFlowable = toMvccEntryFlowable(ctx, commandKeyPredicate);
- return batchOperation(mvccEntryFlowable, ctx, NonBlockingStore::batch);
- }
-
- /**
- * Takes all the modified entries in the flowable and writes or removes them from the stores in a single batch
- * operation per store.
- *
- * The {@link HandleFlowables} is provided for the sole reason of allowing reuse of this method by different callers.
- * @param mvccEntryFlowable flowable containing modified entries
- * @param ctx the context with modifications
- * @param flowableHandler callback handler that actually should subscribe to the underlying store
- * @param key type
- * @param value type
- * @return a stage that when complete will contain how many write operations were done
- */
- private CompletionStage batchOperation(Flowable> mvccEntryFlowable, InvocationContext ctx,
- HandleFlowables flowableHandler) {
- return Single.using(
- this::acquireReadLock,
- ignore -> {
- checkStoreAvailability();
- if (log.isTraceEnabled()) {
- log.trace("Writing batch to stores");
- }
-
- return Flowable.fromIterable(stores)
- .filter(storeStatus -> !storeStatus.characteristics.contains(Characteristic.READ_ONLY))
- .flatMapSingle(storeStatus -> {
- Flowable> flowableToUse;
- boolean shared = storeStatus.config.shared();
- if (shared) {
- if (log.isTraceEnabled()) {
- log.tracef("Store %s is shared, checking skip shared stores and ignoring entries not" +
- " primarily owned by this node", storeStatus.store);
- }
- flowableToUse = mvccEntryFlowable.filter(mvccEntry -> !mvccEntry.isSkipSharedStore());
- } else {
- flowableToUse = mvccEntryFlowable;
- }
-
- boolean segmented = storeStatus.config.segmented();
-
- // Now we have to split this stores' flowable into two (one for remove and one for put)
- flowableToUse = flowableToUse.publish().autoConnect(2);
-
- Flowable> removeFlowable = createRemoveFlowable(
- flowableToUse, shared, segmented, storeStatus);
-
- ByRef.Long writeCount = new ByRef.Long(0);
-
- Flowable>> writeFlowable =
- createWriteFlowable(flowableToUse, ctx, shared, segmented, writeCount, storeStatus);
-
- CompletionStage storeBatchStage = flowableHandler.handleFlowables(storeStatus.store(),
- segmentCount(segmented), removeFlowable, writeFlowable);
-
- return Single.fromCompletionStage(storeBatchStage
- .thenApply(ignore2 -> writeCount.get()));
- // Only take the last element for the count - ensures all stores are completed
- }).last(0L);
-
- },
- this::releaseReadLock
- ).toCompletionStage();
- }
-
- private Flowable> createRemoveFlowable(
- Flowable> flowableToUse, boolean shared, boolean segmented, StoreStatus storeStatus) {
-
- Flowable keyRemoveFlowable = flowableToUse
- .filter(MVCCEntry::isRemoved)
- .map(MVCCEntry::getKey);
-
- Flowable> flowable;
- if (segmented) {
- flowable = keyRemoveFlowable
- .groupBy(keyPartitioner::getSegment)
- .map(SegmentPublisherWrapper::wrap);
- flowable = filterSharedSegments(flowable, null, shared);
- } else {
- if (shared && !isInvalidationCache) {
- keyRemoveFlowable = keyRemoveFlowable.filter(k ->
- distributionManager.getCacheTopology().getDistribution(k).isPrimary());
- }
- flowable = Flowable.just(SingleSegmentPublisher.singleSegment(keyRemoveFlowable));
- }
-
- if (log.isTraceEnabled()) {
- flowable = flowable.doOnSubscribe(sub ->
- log.tracef("Store %s has subscribed to remove batch", storeStatus.store));
- flowable = flowable.map(sp -> {
- int segment = sp.getSegment();
- return SingleSegmentPublisher.singleSegment(segment, Flowable.fromPublisher(sp)
- .doOnNext(keyToRemove -> log.tracef("Emitting key %s for removal from segment %s",
- keyToRemove, segment)));
- });
- }
-
- return flowable;
- }
-
- private Flowable>> createWriteFlowable(
- Flowable> flowableToUse, InvocationContext ctx, boolean shared, boolean segmented,
- ByRef.Long writeCount, StoreStatus storeStatus) {
-
- Flowable> entryWriteFlowable = flowableToUse
- .filter(mvccEntry -> !mvccEntry.isRemoved())
- .map(mvcEntry -> {
- K key = mvcEntry.getKey();
- InternalCacheValue sv = internalEntryFactory.getValueFromCtx(key, ctx);
- //noinspection unchecked
- return (MarshallableEntry) marshallableEntryFactory.create(key, (InternalCacheValue) sv);
- });
-
- Flowable>> flowable;
- if (segmented) {
- // Note the writeCount includes entries that aren't written due to being shared
- // at this point
- entryWriteFlowable = entryWriteFlowable.doOnNext(obj -> writeCount.inc());
- flowable = entryWriteFlowable
- .groupBy(me -> keyPartitioner.getSegment(me.getKey()))
- .map(SegmentPublisherWrapper::wrap);
- // The writeCount will be decremented for each grouping of values ignored
- flowable = filterSharedSegments(flowable, writeCount, shared);
- } else {
- if (shared && !isInvalidationCache) {
- entryWriteFlowable = entryWriteFlowable.filter(me ->
- distributionManager.getCacheTopology().getDistribution(me.getKey()).isPrimary());
- }
- entryWriteFlowable = entryWriteFlowable.doOnNext(obj -> writeCount.inc());
- flowable = Flowable.just(SingleSegmentPublisher.singleSegment(entryWriteFlowable));
- }
-
- if (log.isTraceEnabled()) {
- flowable = flowable.doOnSubscribe(sub ->
- log.tracef("Store %s has subscribed to write batch", storeStatus.store));
- flowable = flowable.map(sp -> {
- int segment = sp.getSegment();
- return SingleSegmentPublisher.singleSegment(segment, Flowable.fromPublisher(sp)
- .doOnNext(me -> log.tracef("Emitting entry %s for write to segment %s",
- me, segment)));
- });
- }
- return flowable;
- }
-
- private Flowable> filterSharedSegments(
- Flowable> flowable, ByRef.Long writeCount, boolean shared) {
- if (!shared || isInvalidationCache) {
- return flowable;
- }
- return flowable.map(sp -> {
- if (distributionManager.getCacheTopology().getSegmentDistribution(sp.getSegment()).isPrimary()) {
- return sp;
- }
- Flowable emptyFlowable = Flowable.fromPublisher(sp);
- if (writeCount != null) {
- emptyFlowable = emptyFlowable.doOnNext(ignore -> writeCount.dec())
- .ignoreElements()
- .toFlowable();
- } else {
- emptyFlowable = emptyFlowable.take(0);
- }
- // Unfortunately we need to still need to subscribe to the publisher even though we don't want
- // the store to use its values. Thus we just return them an empty SegmentPublisher.
- return SingleSegmentPublisher.singleSegment(sp.getSegment(), emptyFlowable);
- });
- }
-
- /**
- * Creates a Flowable of MVCCEntry(s) that were modified due to the commands in the transactional context
- * @param ctx the transactional context
- * @param commandKeyPredicate predicate to test if a key/command combination should be written
- * @param key type
- * @param value type
- * @return a Flowable containing MVCCEntry(s) for the modifications in the tx context
- */
- private Flowable> toMvccEntryFlowable(TxInvocationContext ctx,
- BiPredicate commandKeyPredicate) {
- return Flowable.fromIterable(ctx.getCacheTransaction().getAllModifications())
- .filter(writeCommand -> !writeCommand.hasAnyFlag(FlagBitSets.SKIP_CACHE_STORE | FlagBitSets.ROLLING_UPGRADE))
- .concatMap(writeCommand -> entriesFromCommand(writeCommand, ctx, commandKeyPredicate));
- }
-
- private Flowable> entriesFromCommand(WCT writeCommand, InvocationContext ctx,
- BiPredicate commandKeyPredicate) {
- if (writeCommand instanceof DataWriteCommand) {
- Object key = ((DataWriteCommand) writeCommand).getKey();
- MVCCEntry entry = acquireKeyFromContext(ctx, writeCommand, key, commandKeyPredicate);
- return entry != null ? Flowable.just(entry) : Flowable.empty();
- } else {
- if (writeCommand instanceof InvalidateCommand) {
- return Flowable.empty();
- }
- // Assume multiple key command
- return Flowable.fromIterable(writeCommand.getAffectedKeys())
- .concatMapMaybe(key -> {
- MVCCEntry entry = acquireKeyFromContext(ctx, writeCommand, key, commandKeyPredicate);
- // We use an empty Flowable to symbolize a miss - which is filtered by ofType just below
- return entry != null ? Maybe.just(entry) : Maybe.empty();
- });
- }
- }
-
- private MVCCEntry acquireKeyFromContext(InvocationContext ctx, WCT command, Object key,
- BiPredicate commandKeyPredicate) {
- if (commandKeyPredicate == null || commandKeyPredicate.test(command, key)) {
- //noinspection unchecked
- MVCCEntry entry = (MVCCEntry) ctx.lookupEntry(key);
- if (entry.isChanged()) {
- return entry;
- }
- }
- return null;
- }
-
- /**
- * Here just to create a lambda for method reuse of
- * {@link #batchOperation(Flowable, InvocationContext, HandleFlowables)}
- */
- interface HandleFlowables {
- CompletionStage handleFlowables(NonBlockingStore store, int publisherCount,
- Flowable> removeFlowable,
- Flowable>> putFlowable);
- };
-
- /**
- * Provides a function that groups entries by their segments (via keyPartitioner).
- */
- private Function groupingFunction(Function toKeyFunction) {
- return value -> keyPartitioner.getSegment(toKeyFunction.apply(value));
- }
-
- /**
- * Returns how many segments the user must worry about when segmented or not.
- * @param segmented whether the store is segmented
- * @return how many segments the store must worry about
- */
- private int segmentCount(boolean segmented) {
- return segmented ? segmentCount : 1;
- }
-
- @Override
- public boolean isAvailable() {
- return unavailableExceptionMessage == null;
- }
-
- @Override
- public CompletionStage addSegments(IntSet segments) {
- return Completable.using(
- this::acquireReadLock,
- ignore -> {
- checkStoreAvailability();
- if (log.isTraceEnabled()) {
- log.tracef("Adding segments %s to stores", segments);
- }
- return Flowable.fromIterable(stores)
- .filter(PersistenceManagerImpl::shouldInvokeSegmentMethods)
- .flatMapCompletable(storeStatus -> Completable.fromCompletionStage(storeStatus.store.addSegments(segments)));
- },
- this::releaseReadLock
- ).toCompletionStage(allSegmentedOrShared);
- }
-
- @Override
- public CompletionStage removeSegments(IntSet segments) {
- return Completable.using(
- this::acquireReadLock,
- ignore -> {
- checkStoreAvailability();
- if (log.isTraceEnabled()) {
- log.tracef("Removing segments %s from stores", segments);
- }
- return Flowable.fromIterable(stores)
- .filter(PersistenceManagerImpl::shouldInvokeSegmentMethods)
- .flatMapCompletable(storeStatus -> Completable.fromCompletionStage(storeStatus.store.removeSegments(segments)));
- },
- this::releaseReadLock
- ).toCompletionStage(allSegmentedOrShared);
- }
-
- private static boolean shouldInvokeSegmentMethods(StoreStatus storeStatus) {
- return storeStatus.characteristics.contains(Characteristic.SEGMENTABLE) &&
- !storeStatus.characteristics.contains(Characteristic.SHAREABLE);
- }
-
- public List> getAllStores(Predicate> predicate) {
- long stamp = acquireReadLock();
- try {
- return stores.stream()
- .filter(storeStatus -> predicate.test(storeStatus.characteristics))
- .map(StoreStatus::store)
- .collect(Collectors.toCollection(ArrayList::new));
- } finally {
- releaseReadLock(stamp);
- }
- }
-
- /**
- * Method must be here for augmentation to tell blockhound this method is okay to block
- */
- private long acquireReadLock() {
- return lock.readLock();
- }
-
- /**
- * Opposite of acquireReadLock here for symmetry
- */
- private void releaseReadLock(long stamp) {
- lock.unlockRead(stamp);
- }
-
- private void checkStoreAvailability() {
- if (!enabled) return;
-
- String message = unavailableExceptionMessage;
- if (message != null) {
- throw new StoreUnavailableException(message);
- }
- }
-
- static class StoreStatus {
- final NonBlockingStore store;
- final StoreConfiguration config;
- final Set characteristics;
- // This variable is protected by PersistenceManagerImpl#lock and also the fact that availability check can
- // only be ran one at a time
- boolean availability = true;
-
- StoreStatus(NonBlockingStore store, StoreConfiguration config, Set characteristics) {
- this.store = store;
- this.config = config;
- this.characteristics = characteristics;
- }
-
- NonBlockingStore store() {
- return (NonBlockingStore) store;
- }
- }
-
- boolean anyLocksHeld() {
- return lock.isReadLocked() || lock.isWriteLocked();
- }
-}