package org.jboss.web.loadbalancer.scheduler;
import java.net.MalformedURLException;
import java.net.URL;
import java.util.ArrayList;
import java.util.Iterator;
import javax.management.MBeanServer;
import javax.management.Notification;
import javax.management.NotificationFilter;
import javax.management.NotificationListener;
import javax.management.ObjectName;
import javax.servlet.http.Cookie;
import org.jboss.mx.util.MBeanServerLocator;
import org.jboss.util.xml.XmlHelper;
import org.jboss.web.loadbalancer.util.Constants;
import org.jboss.web.loadbalancer.util.Request;
import org.w3c.dom.Element;
import org.jboss.system.*;
public abstract class AbstractScheduler
extends ServiceMBeanSupport
implements AbstractSchedulerMBean, NotificationListener
{
protected ArrayList hostsUp = new ArrayList();
protected ArrayList hostsDown = new ArrayList();
protected String stickyCookieName;
protected boolean useStickySession = false;
protected Element config;
protected abstract Host getNextHost();
protected void stopService() throws java.lang.Exception
{
deregisterHostMBeans();
}
public void setConfig(Element config)
{
this.config = config;
}
protected void createService() throws java.lang.Exception
{
addHostsFromConfig(config);
if (XmlHelper.getUniqueChildContent(config, "sticky-session").
equalsIgnoreCase("true"))
{
useStickySession = true;
stickyCookieName = XmlHelper.getUniqueChildContent(config,
"sticky-session-cookie-name");
}
if (log.isDebugEnabled())
{
if (useStickySession)
{
log.debug("Using sticky sessions with Cookie name=" + stickyCookieName);
}
else
{
log.debug("Using NO sticky sessions");
}
}
}
protected synchronized void addHostsFromConfig(Element config) throws
Exception
{
Iterator hostIterator = XmlHelper.getChildrenByTagName(XmlHelper.
getUniqueChild(config, "hosts"), "host");
while (hostIterator.hasNext())
{
Element hostElement = (Element) hostIterator.next();
String hostUrl = XmlHelper.getUniqueChildContent(hostElement, "host-url");
if (!hostUrl.endsWith("/"))
{
hostUrl += "/";
}
Host host = addHost(hostUrl);
try
{
host.setLbFactor(Integer.parseInt(XmlHelper.getUniqueChildContent(
hostElement, "lb-factor")));
}
catch (Exception e)
{
log.debug("Could read LbFactor for Host " + host.getUrl() +
" - assuming LbFactor 1");
host.setLbFactor(1);
}
}
}
protected ObjectName genObjectNameForHost(Host host) throws Exception
{
return new ObjectName("jboss.web.loadbalancer: type=Node, protocol=" +
host.getUrl().getProtocol() + ", host=" +
host.getUrl().getHost() + ", port=" +
(host.getUrl().getPort() == -1 ?
host.getUrl().getDefaultPort() :
host.getUrl().getPort()));
}
protected void registerHostMBean(Host host)
{
MBeanServer mbs = MBeanServerLocator.locateJBoss();
try
{
mbs.registerMBean(host, genObjectNameForHost(host));
host.addNotificationListener(this,
new NotificationFilter()
{
public boolean isNotificationEnabled(Notification notification)
{
return (notification instanceof HostStateChangedNotification);
}
}
, host);
}
catch (Exception ex)
{
log.error("Could not register HostMBean", ex);
}
}
protected void deregisterHostMBean(Host host)
{
MBeanServer mbs = MBeanServerLocator.locateJBoss();
try
{
mbs.unregisterMBean(genObjectNameForHost(host));
host.removeNotificationListener(this);
}
catch (Exception ex)
{
log.error("Could not unregister HostMBean", ex);
}
}
protected synchronized void deregisterHostMBeans() throws Exception
{
Iterator hostIterator = hostsUp.iterator();
while (hostIterator.hasNext())
{
Host host = (Host) hostIterator.next();
deregisterHostMBean(host);
}
hostIterator = hostsDown.iterator();
while (hostIterator.hasNext())
{
Host host = (Host) hostIterator.next();
deregisterHostMBean(host);
}
}
public void getHost(Request schedRequest) throws
NoHostAvailableException
{
Host host = null;
host = findStickyHost(schedRequest);
while (host == null)
{
if (hostsUp.size() == 0)
{
throw new NoHostAvailableException("No host to schedule request");
}
host = getNextHost();
}
schedRequest.setHost(host);
if (useStickySession)
{
setStickyCookie(schedRequest);
}
}
protected void setStickyCookie(Request schedRequest)
{
Cookie cookie = new Cookie(stickyCookieName,
Integer.toString(schedRequest.getHost().hashCode()));
cookie.setPath("/");
cookie.setMaxAge( -1);
schedRequest.getResponse().addCookie(cookie);
}
protected Host findStickyHost(Request schedRequest)
{
Host host = null;
if (useStickySession)
{
Cookie[] cookies = schedRequest.getRequest().getCookies();
for (int i = 0; cookies != null && i < cookies.length; ++i)
{
Cookie cookie = cookies[i];
if (cookie.getName().equals(stickyCookieName))
{
log.debug("Sticky Cookie found!");
int cookieHash = Integer.parseInt(cookie.getValue());
Iterator iter = hostsUp.iterator();
while (iter.hasNext())
{
Host tempHost = (Host) iter.next();
if (tempHost.hashCode() == cookieHash)
{
host = tempHost;
if (log.isDebugEnabled())
{
log.debug("Sticky Cookie sticks client to host with URL " +
tempHost.toString());
}
break;
}
}
break;
}
}
return host;
}
else
{
return null;
}
}
protected synchronized void markNodeDown(Host host)
{
int index = hostsUp.indexOf(host);
if (index == -1)
{
return;
}
hostsUp.remove(index);
hostsDown.add(host);
}
protected synchronized void markNodeUp(Host host)
{
int index = hostsDown.indexOf(host);
if (index == -1)
{
return;
}
hostsDown.remove(index);
hostsUp.add(host);
}
public ArrayList getHostsDown()
{
return hostsDown;
}
public ArrayList getHostsUp()
{
return hostsUp;
}
public synchronized Host addHost(String hostString) throws
MalformedURLException
{
if (hostString == null)
{
return null;
}
if (!hostString.endsWith("/"))
{
hostString += "/";
}
Host host = new Host(new URL(hostString));
if (hostsUp.indexOf(host) > -1 || hostsDown.indexOf(host) > -1)
{
log.info("Host " + hostString + " already there. Ignored");
return null;
}
hostsUp.add(host);
registerHostMBean(host);
return host;
}
public synchronized void removeHost(URL url)
{
if (url == null)
{
return;
}
Host host = new Host(url);
int index = hostsUp.indexOf(host);
if (index > -1)
{
deregisterHostMBean( (Host) hostsUp.get(index));
hostsUp.remove(index);
}
index = hostsDown.indexOf(host);
if (index > -1)
{
deregisterHostMBean( (Host) hostsDown.get(index));
hostsDown.remove(index);
}
}
public void handleNotification(Notification notification, Object handback)
{
HostStateChangedNotification hscn=(HostStateChangedNotification)notification;
Host host=(Host)hscn.getSource();
switch (host.getState())
{
case Constants.STATE_NODE_UP:
markNodeUp(host);
break;
case Constants.STATE_NODE_DOWN:
case Constants.STATE_NODE_FORCED_DOWN:
markNodeDown(host);
break;
}
}
}