/**
 * 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.transport;

import java.net.Socket;
import java.net.SocketException;
import java.net.URI;
import java.util.concurrent.TimeUnit;

import javax.jms.Connection;
import javax.jms.Destination;
import javax.jms.JMSException;
import javax.jms.MessageConsumer;
import javax.jms.Session;

import junit.framework.Test;

import org.apache.activemq.ActiveMQConnectionFactory;
import org.apache.activemq.JmsTestSupport;
import org.apache.activemq.broker.BrokerService;
import org.apache.activemq.command.ActiveMQQueue;
import org.apache.activemq.store.kahadb.KahaDBPersistenceAdapter;
import org.apache.activemq.transport.stomp.Stomp;
import org.apache.activemq.transport.stomp.StompConnection;
import org.apache.activemq.util.SocketProxy;
import org.apache.activemq.util.URISupport;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

public class SoWriteTimeoutTest extends JmsTestSupport {

   private static final Logger LOG = LoggerFactory.getLogger(SoWriteTimeoutTest.class);

   final int receiveBufferSize = 16 * 1024;
   public String brokerTransportScheme = "nio";

   protected BrokerService createBroker() throws Exception {
      BrokerService broker = super.createBroker();
      broker.setPersistent(true);
      broker.setDeleteAllMessagesOnStartup(true);
      KahaDBPersistenceAdapter adapter = new KahaDBPersistenceAdapter();
      adapter.setConcurrentStoreAndDispatchQueues(false);
      broker.setPersistenceAdapter(adapter);
      broker.addConnector(brokerTransportScheme + "://localhost:0?wireFormat.maxInactivityDuration=0&transport.soWriteTimeout=1000&transport.sleep=1000");
      if ("nio".equals(brokerTransportScheme)) {
         broker.addConnector("stomp+" + brokerTransportScheme + "://localhost:0?transport.soWriteTimeout=1000&transport.sleep=1000&socketBufferSize=" + receiveBufferSize + "&trace=true");
      }
      return broker;
   }

   public void initCombosForTestWriteTimeout() {
      addCombinationValues("brokerTransportScheme", new Object[]{"tcp", "nio"});
   }

   public void testWriteTimeout() throws Exception {

      Destination dest = new ActiveMQQueue("testWriteTimeout");
      messageTextPrefix = initMessagePrefix(8 * 1024);
      sendMessages(dest, 500);

      URI tcpBrokerUri = URISupport.removeQuery(broker.getTransportConnectors().get(0).getConnectUri());
      LOG.info("consuming using uri: " + tcpBrokerUri);

      SocketProxy proxy = new SocketProxy();
      proxy.setTarget(tcpBrokerUri);
      proxy.setReceiveBufferSize(receiveBufferSize);
      proxy.open();

      ActiveMQConnectionFactory factory = new ActiveMQConnectionFactory(proxy.getUrl());
      Connection c = factory.createConnection();
      c.start();
      Session session = c.createSession(true, Session.SESSION_TRANSACTED);
      MessageConsumer consumer = session.createConsumer(dest);
      proxy.pause();
      // writes should back up... writeTimeout will kick in and abort the connection
      TimeUnit.SECONDS.sleep(10);
      proxy.goOn();
      assertNotNull("can receive buffered messages", consumer.receive(500));
      try {
         session.commit();
         fail("expect commit to fail as server has aborted writeTimeout connection");
      }
      catch (JMSException expected) {
      }
   }

   public void testWriteTimeoutStompNio() throws Exception {
      ActiveMQQueue dest = new ActiveMQQueue("testWriteTimeout");
      messageTextPrefix = initMessagePrefix(8 * 1024);
      sendMessages(dest, 500);

      URI stompBrokerUri = URISupport.removeQuery(broker.getTransportConnectors().get(1).getConnectUri());
      LOG.info("consuming using uri: " + stompBrokerUri);

      SocketProxy proxy = new SocketProxy();
      proxy.setTarget(new URI("tcp://localhost:" + stompBrokerUri.getPort()));
      proxy.setReceiveBufferSize(receiveBufferSize);
      proxy.open();

      StompConnection stompConnection = new StompConnection();
      stompConnection.open(new Socket("localhost", proxy.getUrl().getPort()));
      stompConnection.getStompSocket().setTcpNoDelay(true);

      String frame = "CONNECT\n" + "login:system\n" + "passcode:manager\n\n" + Stomp.NULL;
      stompConnection.sendFrame(frame);
      frame = stompConnection.receiveFrame();
      assertTrue(frame.startsWith("CONNECTED"));

      frame = "SUBSCRIBE\n" + "destination:/queue/" + dest.getQueueName() + "\n" + "ack:client\n\n" + Stomp.NULL;
      stompConnection.sendFrame(frame);

      // ensure dispatch has started before pause
      frame = stompConnection.receiveFrame();
      assertTrue(frame.startsWith("MESSAGE"));

      proxy.pause();

      // writes should back up... writeTimeout will kick in and abort the connection
      TimeUnit.SECONDS.sleep(1);

      // see the blocked threads
      //dumpAllThreads("blocked on write");

      // abort should be done after this
      TimeUnit.SECONDS.sleep(10);

      proxy.goOn();

      // get a buffered message
      frame = stompConnection.receiveFrame();
      assertTrue(frame.startsWith("MESSAGE"));

      // verify connection is dead
      try {
         for (int i = 0; i < 200; i++) {
            stompConnection.send("/queue/" + dest.getPhysicalName(), "ShouldBeDeadConnectionText" + i);
         }
         fail("expected send to fail with timeout out connection");
      }
      catch (SocketException expected) {
         LOG.info("got exception on send after timeout: " + expected);
      }
   }

   private String initMessagePrefix(int i) {
      byte[] content = new byte[i];
      return new String(content);
   }

   @Override
   protected void setUp() throws Exception {
      setAutoFail(true);
      super.setUp();
   }

   public static Test suite() {
      return suite(SoWriteTimeoutTest.class);
   }
}
