/**
* Licensed to the Apache Software Foundation (ASF) under one or more
* contributor license agreements. See the NOTICE file distributed with
* this work for additional information regarding copyright ownership.
* The ASF licenses this file to You under the Apache License, Version 2.0
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package org.apache.activemq.apollo.openwire.support.state;
import java.io.IOException;
import java.util.Iterator;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import org.apache.activemq.apollo.util.IOExceptionSupport;
import org.apache.activemq.apollo.util.VoidFunction1;
import org.apache.activemq.apollo.openwire.command.Command;
import org.apache.activemq.apollo.openwire.command.ConnectionId;
import org.apache.activemq.apollo.openwire.command.ConnectionInfo;
import org.apache.activemq.apollo.openwire.command.ConsumerId;
import org.apache.activemq.apollo.openwire.command.ConsumerInfo;
import org.apache.activemq.apollo.openwire.command.DestinationInfo;
import org.apache.activemq.apollo.openwire.command.Message;
import org.apache.activemq.apollo.openwire.command.MessageId;
import org.apache.activemq.apollo.openwire.command.ProducerId;
import org.apache.activemq.apollo.openwire.command.ProducerInfo;
import org.apache.activemq.apollo.openwire.command.Response;
import org.apache.activemq.apollo.openwire.command.SessionId;
import org.apache.activemq.apollo.openwire.command.SessionInfo;
import org.apache.activemq.apollo.openwire.command.TransactionInfo;
import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;
/**
* Tracks the state of a connection so a newly established transport can be
* re-initialized to the state that was tracked.
*
*/
public class ConnectionStateTracker extends CommandVisitorAdapter {
private static final Log LOG = LogFactory.getLog(ConnectionStateTracker.class);
private static final org.apache.activemq.apollo.openwire.support.state.Tracked TRACKED_RESPONSE_MARKER = new org.apache.activemq.apollo.openwire.support.state.Tracked(null);
protected final ConcurrentHashMap<ConnectionId, org.apache.activemq.apollo.openwire.support.state.ConnectionState> connectionStates = new ConcurrentHashMap<ConnectionId, org.apache.activemq.apollo.openwire.support.state.ConnectionState>();
private boolean trackTransactions;
private boolean restoreSessions = true;
private boolean restoreConsumers = true;
private boolean restoreProducers = true;
private boolean restoreTransaction = true;
private boolean trackMessages = true;
private int maxCacheSize = 128 * 1024;
private int currentCacheSize;
private Map<MessageId,Message> messageCache = new LinkedHashMap<MessageId,Message>(){
protected boolean removeEldestEntry(Map.Entry<MessageId,Message> eldest) {
boolean result = currentCacheSize > maxCacheSize;
if (result) {
currentCacheSize -= eldest.getValue().getSize();
}
return result;
}
};
private class RemoveTransactionAction implements Runnable {
private final TransactionInfo info;
public RemoveTransactionAction(TransactionInfo info) {
this.info = info;
}
public void run() {
ConnectionId connectionId = info.getConnectionId();
org.apache.activemq.apollo.openwire.support.state.ConnectionState cs = connectionStates.get(connectionId);
cs.removeTransactionState(info.getTransactionId());
}
}
/**
*
*
* @param command
* @return null if the command is not state tracked.
* @throws IOException
*/
public org.apache.activemq.apollo.openwire.support.state.Tracked track(Command command) throws IOException {
try {
return (org.apache.activemq.apollo.openwire.support.state.Tracked)command.visit(this);
} catch (IOException e) {
throw e;
} catch (Throwable e) {
throw IOExceptionSupport.create(e);
}
}
public void trackBack(Command command) {
if (trackMessages && command != null && command.isMessage()) {
Message message = (Message) command;
if (message.getTransactionId()==null) {
currentCacheSize = currentCacheSize + message.getSize();
}
}
}
public void restore(VoidFunction1<Command> transport) throws IOException {
// Restore the connections.
for (Iterator<org.apache.activemq.apollo.openwire.support.state.ConnectionState> iter = connectionStates.values().iterator(); iter.hasNext();) {
org.apache.activemq.apollo.openwire.support.state.ConnectionState connectionState = iter.next();
transport.apply(connectionState.getInfo());
restoreTempDestinations(transport, connectionState);
if (restoreSessions) {
restoreSessions(transport, connectionState);
}
if (restoreTransaction) {
restoreTransactions(transport, connectionState);
}
}
//now flush messages
for (Message msg:messageCache.values()) {
transport.apply(msg);
}
}
private void restoreTransactions(VoidFunction1<Command> transport, org.apache.activemq.apollo.openwire.support.state.ConnectionState connectionState) throws IOException {
for (Iterator iter = connectionState.getTransactionStates().iterator(); iter.hasNext();) {
TransactionState transactionState = (TransactionState)iter.next();
if (LOG.isDebugEnabled()) {
LOG.debug("tx: " + transactionState.getId());
}
for (Iterator iterator = transactionState.getCommands().iterator(); iterator.hasNext();) {
Command command = (Command)iterator.next();
if (LOG.isDebugEnabled()) {
LOG.debug("tx replay: " + command);
}
transport.apply(command);
}
}
}
/**
* @param transport
* @param connectionState
* @throws IOException
*/
protected void restoreSessions(VoidFunction1<Command> transport, org.apache.activemq.apollo.openwire.support.state.ConnectionState connectionState) throws IOException {
// Restore the connection's sessions
for (Iterator iter2 = connectionState.getSessionStates().iterator(); iter2.hasNext();) {
SessionState sessionState = (SessionState)iter2.next();
transport.apply(sessionState.getInfo());
if (restoreProducers) {
restoreProducers(transport, sessionState);
}
if (restoreConsumers) {
restoreConsumers(transport, sessionState);
}
}
}
/**
* @param transport
* @param sessionState
* @throws IOException
*/
protected void restoreConsumers(VoidFunction1<Command> transport, SessionState sessionState) throws IOException {
// Restore the session's consumers
for (Iterator iter3 = sessionState.getConsumerStates().iterator(); iter3.hasNext();) {
org.apache.activemq.apollo.openwire.support.state.ConsumerState consumerState = (org.apache.activemq.apollo.openwire.support.state.ConsumerState)iter3.next();
transport.apply(consumerState.getInfo());
}
}
/**
* @param transport
* @param sessionState
* @throws IOException
*/
protected void restoreProducers(VoidFunction1<Command> transport, SessionState sessionState) throws IOException {
// Restore the session's producers
for (Iterator iter3 = sessionState.getProducerStates().iterator(); iter3.hasNext();) {
org.apache.activemq.apollo.openwire.support.state.ProducerState producerState = (org.apache.activemq.apollo.openwire.support.state.ProducerState)iter3.next();
transport.apply(producerState.getInfo());
}
}
/**
* @param transport
* @param connectionState
* @throws IOException
*/
protected void restoreTempDestinations(VoidFunction1<Command> transport, org.apache.activemq.apollo.openwire.support.state.ConnectionState connectionState)
throws IOException {
// Restore the connection's temp destinations.
for (Iterator iter2 = connectionState.getTempDesinations().iterator(); iter2.hasNext();) {
transport.apply((DestinationInfo)iter2.next());
}
}
public Response processAddDestination(DestinationInfo info) {
if (info != null) {
org.apache.activemq.apollo.openwire.support.state.ConnectionState cs = connectionStates.get(info.getConnectionId());
if (cs != null && info.getDestination().isTemporary()) {
cs.addTempDestination(info);
}
}
return TRACKED_RESPONSE_MARKER;
}
public Response processRemoveDestination(DestinationInfo info) {
if (info != null) {
org.apache.activemq.apollo.openwire.support.state.ConnectionState cs = connectionStates.get(info.getConnectionId());
if (cs != null && info.getDestination().isTemporary()) {
cs.removeTempDestination(info.getDestination());
}
}
return TRACKED_RESPONSE_MARKER;
}
public Response processAddProducer(ProducerInfo info) {
if (info != null && info.getProducerId() != null) {
SessionId sessionId = info.getProducerId().getParentId();
if (sessionId != null) {
ConnectionId connectionId = sessionId.getParentId();
if (connectionId != null) {
org.apache.activemq.apollo.openwire.support.state.ConnectionState cs = connectionStates.get(connectionId);
if (cs != null) {
SessionState ss = cs.getSessionState(sessionId);
if (ss != null) {
ss.addProducer(info);
}
}
}
}
}
return TRACKED_RESPONSE_MARKER;
}
public Response processRemoveProducer(ProducerId id) {
if (id != null) {
SessionId sessionId = id.getParentId();
if (sessionId != null) {
ConnectionId connectionId = sessionId.getParentId();
if (connectionId != null) {
org.apache.activemq.apollo.openwire.support.state.ConnectionState cs = connectionStates.get(connectionId);
if (cs != null) {
SessionState ss = cs.getSessionState(sessionId);
if (ss != null) {
ss.removeProducer(id);
}
}
}
}
}
return TRACKED_RESPONSE_MARKER;
}
public Response processAddConsumer(ConsumerInfo info) {
if (info != null) {
SessionId sessionId = info.getConsumerId().getParentId();
if (sessionId != null) {
ConnectionId connectionId = sessionId.getParentId();
if (connectionId != null) {
org.apache.activemq.apollo.openwire.support.state.ConnectionState cs = connectionStates.get(connectionId);
if (cs != null) {
SessionState ss = cs.getSessionState(sessionId);
if (ss != null) {
ss.addConsumer(info);
}
}
}
}
}
return TRACKED_RESPONSE_MARKER;
}
public Response processRemoveConsumer(ConsumerId id, long lastDeliveredSequenceId) {
if (id != null) {
SessionId sessionId = id.getParentId();
if (sessionId != null) {
ConnectionId connectionId = sessionId.getParentId();
if (connectionId != null) {
org.apache.activemq.apollo.openwire.support.state.ConnectionState cs = connectionStates.get(connectionId);
if (cs != null) {
SessionState ss = cs.getSessionState(sessionId);
if (ss != null) {
ss.removeConsumer(id);
}
}
}
}
}
return TRACKED_RESPONSE_MARKER;
}
public Response processAddSession(SessionInfo info) {
if (info != null) {
ConnectionId connectionId = info.getSessionId().getParentId();
if (connectionId != null) {
org.apache.activemq.apollo.openwire.support.state.ConnectionState cs = connectionStates.get(connectionId);
if (cs != null) {
cs.addSession(info);
}
}
}
return TRACKED_RESPONSE_MARKER;
}
public Response processRemoveSession(SessionId id, long lastDeliveredSequenceId) {
if (id != null) {
ConnectionId connectionId = id.getParentId();
if (connectionId != null) {
org.apache.activemq.apollo.openwire.support.state.ConnectionState cs = connectionStates.get(connectionId);
if (cs != null) {
cs.removeSession(id);
}
}
}
return TRACKED_RESPONSE_MARKER;
}
public Response processAddConnection(ConnectionInfo info) {
if (info != null) {
connectionStates.put(info.getConnectionId(), new org.apache.activemq.apollo.openwire.support.state.ConnectionState(info));
}
return TRACKED_RESPONSE_MARKER;
}
public Response processRemoveConnection(ConnectionId id, long lastDeliveredSequenceId) throws Exception {
if (id != null) {
connectionStates.remove(id);
}
return TRACKED_RESPONSE_MARKER;
}
public Response processMessage(Message send) throws Exception {
if (send != null) {
if (trackTransactions && send.getTransactionId() != null) {
ConnectionId connectionId = send.getProducerId().getParentId().getParentId();
if (connectionId != null) {
org.apache.activemq.apollo.openwire.support.state.ConnectionState cs = connectionStates.get(connectionId);
if (cs != null) {
TransactionState transactionState = cs.getTransactionState(send.getTransactionId());
if (transactionState != null) {
transactionState.addCommand(send);
}
}
}
return TRACKED_RESPONSE_MARKER;
}else if (trackMessages) {
messageCache.put(send.getMessageId(), send.copy());
}
}
return null;
}
public Response processBeginTransaction(TransactionInfo info) {
if (trackTransactions && info != null && info.getTransactionId() != null) {
ConnectionId connectionId = info.getConnectionId();
if (connectionId != null) {
org.apache.activemq.apollo.openwire.support.state.ConnectionState cs = connectionStates.get(connectionId);
if (cs != null) {
cs.addTransactionState(info.getTransactionId());
TransactionState state = cs.getTransactionState(info.getTransactionId());
state.addCommand(info);
}
}
return TRACKED_RESPONSE_MARKER;
}
return null;
}
public Response processPrepareTransaction(TransactionInfo info) throws Exception {
if (trackTransactions && info != null) {
ConnectionId connectionId = info.getConnectionId();
if (connectionId != null) {
org.apache.activemq.apollo.openwire.support.state.ConnectionState cs = connectionStates.get(connectionId);
if (cs != null) {
TransactionState transactionState = cs.getTransactionState(info.getTransactionId());
if (transactionState != null) {
transactionState.addCommand(info);
}
}
}
return TRACKED_RESPONSE_MARKER;
}
return null;
}
public Response processCommitTransactionOnePhase(TransactionInfo info) throws Exception {
if (trackTransactions && info != null) {
ConnectionId connectionId = info.getConnectionId();
if (connectionId != null) {
org.apache.activemq.apollo.openwire.support.state.ConnectionState cs = connectionStates.get(connectionId);
if (cs != null) {
TransactionState transactionState = cs.getTransactionState(info.getTransactionId());
if (transactionState != null) {
transactionState.addCommand(info);
return new org.apache.activemq.apollo.openwire.support.state.Tracked(new RemoveTransactionAction(info));
}
}
}
}
return null;
}
public Response processCommitTransactionTwoPhase(TransactionInfo info) throws Exception {
if (trackTransactions && info != null) {
ConnectionId connectionId = info.getConnectionId();
if (connectionId != null) {
org.apache.activemq.apollo.openwire.support.state.ConnectionState cs = connectionStates.get(connectionId);
if (cs != null) {
TransactionState transactionState = cs.getTransactionState(info.getTransactionId());
if (transactionState != null) {
transactionState.addCommand(info);
return new org.apache.activemq.apollo.openwire.support.state.Tracked(new RemoveTransactionAction(info));
}
}
}
}
return null;
}
public Response processRollbackTransaction(TransactionInfo info) throws Exception {
if (trackTransactions && info != null) {
ConnectionId connectionId = info.getConnectionId();
if (connectionId != null) {
org.apache.activemq.apollo.openwire.support.state.ConnectionState cs = connectionStates.get(connectionId);
if (cs != null) {
TransactionState transactionState = cs.getTransactionState(info.getTransactionId());
if (transactionState != null) {
transactionState.addCommand(info);
return new org.apache.activemq.apollo.openwire.support.state.Tracked(new RemoveTransactionAction(info));
}
}
}
}
return null;
}
public Response processEndTransaction(TransactionInfo info) throws Exception {
if (trackTransactions && info != null) {
ConnectionId connectionId = info.getConnectionId();
if (connectionId != null) {
org.apache.activemq.apollo.openwire.support.state.ConnectionState cs = connectionStates.get(connectionId);
if (cs != null) {
TransactionState transactionState = cs.getTransactionState(info.getTransactionId());
if (transactionState != null) {
transactionState.addCommand(info);
}
}
}
return TRACKED_RESPONSE_MARKER;
}
return null;
}
public boolean isRestoreConsumers() {
return restoreConsumers;
}
public void setRestoreConsumers(boolean restoreConsumers) {
this.restoreConsumers = restoreConsumers;
}
public boolean isRestoreProducers() {
return restoreProducers;
}
public void setRestoreProducers(boolean restoreProducers) {
this.restoreProducers = restoreProducers;
}
public boolean isRestoreSessions() {
return restoreSessions;
}
public void setRestoreSessions(boolean restoreSessions) {
this.restoreSessions = restoreSessions;
}
public boolean isTrackTransactions() {
return trackTransactions;
}
public void setTrackTransactions(boolean trackTransactions) {
this.trackTransactions = trackTransactions;
}
public boolean isRestoreTransaction() {
return restoreTransaction;
}
public void setRestoreTransaction(boolean restoreTransaction) {
this.restoreTransaction = restoreTransaction;
}
public boolean isTrackMessages() {
return trackMessages;
}
public void setTrackMessages(boolean trackMessages) {
this.trackMessages = trackMessages;
}
public int getMaxCacheSize() {
return maxCacheSize;
}
public void setMaxCacheSize(int maxCacheSize) {
this.maxCacheSize = maxCacheSize;
}
}