- public Short getTableIdForTxId(NodeId nodeId,TransactionId id){
- return txIdTotableIdMap.get(new TxIdEntry(nodeId,id,null));
+ public void recordExpectedTableTransaction(TransactionId id, Short tableId) {
+ recordExpectedTransaction(id);
+ txIdTotableIdMap.put(new TxIdEntry(id), Preconditions.checkNotNull(tableId));
+ }
+
+ public Short isExpectedTableTransaction(TransactionAware transaction) {
+ Boolean more = null;
+ if (transaction instanceof MultipartTransactionAware) {
+ more = ((MultipartTransactionAware)transaction).isMoreReplies();
+ }
+
+ if (!isExpectedTransaction(transaction, more)) {
+ return null;
+ }
+
+ final TxIdEntry key = new TxIdEntry(transaction.getTransactionId());
+ if (more != null && more.booleanValue()) {
+ return txIdTotableIdMap.get(key);
+ } else {
+ return txIdTotableIdMap.remove(key);
+ }