- private static Map<TransactionId,Short> txIdTotableIdMap = new ConcurrentHashMap<TransactionId,Short>();
-
- public MultipartMessageManager(){}
-
- public Short getTableIdForTxId(TransactionId id){
-
- return txIdTotableIdMap.get(id);
-
+ private final Map<TxIdEntry,Short> txIdTotableIdMap = new ConcurrentHashMap<>();
+
+ private static final class TxIdEntry {
+ private final StatsRequestType requestType;
+ private final TransactionId txId;
+
+ public TxIdEntry(TransactionId txId, StatsRequestType requestType){
+ this.txId = txId;
+ this.requestType = requestType;
+ }
+ public TransactionId getTxId() {
+ return txId;
+ }
+ public StatsRequestType getRequestType() {
+ return requestType;
+ }
+ @Override
+ public int hashCode() {
+ final int prime = 31;
+ int result = 1;
+ result = prime * result + ((txId == null) ? 0 : txId.hashCode());
+ return result;
+ }
+ @Override
+ public boolean equals(Object obj) {
+ if (this == obj) {
+ return true;
+ }
+ if (obj == null) {
+ return false;
+ }
+ if (!(obj instanceof TxIdEntry)) {
+ return false;
+ }
+ TxIdEntry other = (TxIdEntry) obj;
+
+ if (txId == null) {
+ if (other.txId != null) {
+ return false;
+ }
+ } else if (!txId.equals(other.txId)) {
+ return false;
+ }
+ return true;
+ }
+
+ @Override
+ public String toString() {
+ return "TxIdEntry [txId=" + txId + ", requestType=" + requestType + "]";
+ }
+ }
+
+ public void recordExpectedTableTransaction(TransactionId id, StatsRequestType type, Short tableId) {
+ recordExpectedTransaction(id, type);
+ txIdTotableIdMap.put(new TxIdEntry(id, null), Preconditions.checkNotNull(tableId));
+ }
+
+ public Short isExpectedTableTransaction(TransactionAware transaction, Boolean more) {
+ if (!isExpectedTransaction(transaction, more)) {
+ return null;
+ }
+
+ final TxIdEntry key = new TxIdEntry(transaction.getTransactionId(), null);
+ if (more != null && more.booleanValue()) {
+ return txIdTotableIdMap.get(key);
+ } else {
+ return txIdTotableIdMap.remove(key);
+ }