// $Id: FC.java.txt,v 1.2 2008/01/22 10:44:29 belaban Exp $

package org.jgroups.protocols;

import EDU.oswego.cs.dl.util.concurrent.ConcurrentReaderHashMap;
import org.jgroups.*;
import org.jgroups.stack.Protocol;
import org.jgroups.util.BoundedList;
import org.jgroups.util.CondVar;
import org.jgroups.util.Streamable;

import java.io.*;
import java.util.*;

/**
 * Simple flow control protocol based on a credit system. Each sender has a number of credits (bytes
 * to send). When the credits have been exhausted, the sender blocks. Each receiver also keeps track of
 * how many credits it has received from a sender. When credits for a sender fall below a threshold,
 * the receiver sends more credits to the sender. Works for both unicast and multicast messages.
 * <p>
 * Note that this protocol must be located towards the top of the stack, or all down_threads from JChannel to this
 * protocol must be set to false ! This is in order to block JChannel.send()/JChannel.down().
 * @author Bela Ban
 * @version $Revision: 1.2 $
 */
public class FC extends Protocol {

    /** My own address */
    Address local_addr=null;

    /** HashMap<Address,Long>: keys are members, values are credits left. For each send, the
     * number of credits is decremented by the message size */
    final Map sent=new HashMap(11);
    // final Map sent=new ConcurrentHashMap(11);

    /** HashMap<Address,Long>: keys are members, values are credits left (in bytes).
     * For each receive, the credits for the sender are decremented by the size of the received message.
     * When the credits are 0, we refill and send a CREDIT message to the sender. Sender blocks until CREDIT
     * is received after reaching <tt>min_credits</tt> credits. */
    final Map received=new ConcurrentReaderHashMap(11);
    // final Map received=new ConcurrentHashMap(11);

    /** We cache the membership */
    final Vector members=new Vector(11);

    /** List of members from whom we expect credits */
    final Vector creditors=new Vector(11);

    /** Max number of bytes to send per receiver until an ack must
     * be received before continuing sending */
    private long max_credits=50000;

    /** Max time (in milliseconds) to block. If credit hasn't been received after max_block_time, we send
     * a REPLENISHMENT request to the members from which we expect credits. A value <= 0 means to
     * wait forever.
     */
    private long max_block_time=5000;

    /** If credits fall below this limit, we send more credits to the sender. (We also send when
     * credits are exhausted (0 credits left)) */
    double min_threshold=0.25;

    /** Computed as <tt>max_credits</tt> times <tt>min_theshold</tt>. If explicitly set, this will
     * override the above computation */
    private long min_credits=0;

    /** Current blocking. True if blocking, else false */
    private final CondVar blocking=new CondVar("blocking", Boolean.FALSE);

    static final String name="FC";

    private long start_blocking=0, stop_blocking=0;

    private int num_blockings=0, num_replenishments=0, num_credit_requests=0;
    private long total_time_blocking=0;

    final BoundedList last_blockings=new BoundedList(50);

    final static FcHeader REPLENISH_HDR=new FcHeader(FcHeader.REPLENISH);
    final static FcHeader CREDIT_REQUEST_HDR=new FcHeader(FcHeader.CREDIT_REQUEST);



    public String getName() {
        return name;
    }

    public void resetStats() {
        super.resetStats();
        num_blockings=num_replenishments=num_credit_requests=0;
        total_time_blocking=0;
        last_blockings.removeAll();
    }

    public long getMaxCredits() {
        return max_credits;
    }

    public void setMaxCredits(long max_credits) {
        this.max_credits=max_credits;
    }

    public double getMinThreshold() {
        return min_threshold;
    }

    public void setMinThreshold(double min_threshold) {
        this.min_threshold=min_threshold;
    }

    public long getMinCredits() {
        return min_credits;
    }

    public void setMinCredits(long min_credits) {
        this.min_credits=min_credits;
    }

    public boolean isBlocked() {
        Object obj=blocking.get();
        return obj != null && obj instanceof Boolean && ((Boolean)obj).booleanValue();
    }

    public int getNumberOfBlockings() {
        return num_blockings;
    }

    public long getTotalTimeBlocked() {
        return total_time_blocking;
    }

    public double getAverageTimeBlocked() {
        return num_blockings == 0? num_blockings : total_time_blocking / num_blockings;
    }

    public int getNumberOfReplenishmentsReceived() {
        return num_replenishments;
    }

    public int getNumberOfCreditRequests() {
        return num_credit_requests;
    }

    public String printSenderCredits() {
        return printMap(sent);
    }

    public String printReceiverCredits() {
        return printMap(received);
    }

    public String printCredits() {
        StringBuilder sb=new StringBuilder();
        sb.append("senders:\n").append(printMap(sent)).append("\n\nreceivers:\n").append(printMap(received));
        return sb.toString();
    }

    public Map dumpStats() {
        Map retval=super.dumpStats();
        if(retval == null)
            retval=new HashMap();
        retval.put("senders", printMap(sent));
        retval.put("receivers", printMap(received));
        retval.put("num_blockings", new Integer(this.num_blockings));
        retval.put("avg_time_blocked", new Double(getAverageTimeBlocked()));
        retval.put("num_replenishments", new Integer(this.num_replenishments));
        return retval;
    }

    public String showLastBlockingTimes() {
        return last_blockings.toString();
    }



    public void unblock() {
        unblockSender();
    }



    public boolean setProperties(Properties props) {
        String  str;
        boolean min_credits_set=false;

        super.setProperties(props);
        str=props.getProperty("max_credits");
        if(str != null) {
            max_credits=Long.parseLong(str);
            props.remove("max_credits");
        }

        str=props.getProperty("min_threshold");
        if(str != null) {
            min_threshold=Double.parseDouble(str);
            props.remove("min_threshold");
        }

        str=props.getProperty("min_credits");
        if(str != null) {
            min_credits=Long.parseLong(str);
            props.remove("min_credits");
            min_credits_set=true;
        }

        if(!min_credits_set)
            min_credits=(long)((double)max_credits * min_threshold);

        str=props.getProperty("max_block_time");
        if(str != null) {
            max_block_time=Long.parseLong(str);
            props.remove("max_block_time");
        }

        if(props.size() > 0) {
            log.error("FC.setProperties(): the following properties are not recognized: " + props);

            return false;
        }
        return true;
    }

    public void stop() {
        super.stop();
        unblock();
    }


    /**
     * We need to receive view changes concurrent to messages on the down events: a message might blocks, e.g.
     * because we don't have enough credits to send to member P. However, if member P crashed, we need to unblock !
     * @param evt
     */
    protected void receiveDownEvent(Event evt) {
        if(evt.getType() == Event.VIEW_CHANGE) {
            View v=(View)evt.getArg();
            Vector mbrs=v.getMembers();
            handleViewChange(mbrs);
        }
        super.receiveDownEvent(evt);
    }

    public void down(Event evt) {
        switch(evt.getType()) {
        case Event.MSG:
            handleDownMessage(evt);
            return;
        }
        passDown(evt); // this could potentially use the lower protocol's thread which may block
    }


    private synchronized void handleDownMessage(Event evt) {
        if(Boolean.TRUE.equals(blocking.get())) { // blocked
            waitUntilEnoughCreditsAvailable();
        }
        else {
            // not blocked
            boolean rc;
            synchronized(sent) { // 'sent' is the same lock as blocking.getLock()...
                rc=decrMessage((Message)evt.getArg());
                if(rc == false) {
                    if(trace)
                        log.trace("blocking due to insufficient credits");
                    blocking.set(Boolean.TRUE);
                    start_blocking=System.currentTimeMillis();
                    num_blockings++;
                }
            }
            if(rc == false) {
                waitUntilEnoughCreditsAvailable();
            }
        }

        passDown(evt);
    }



    public void up(Event evt) {
        switch(evt.getType()) {
            case Event.SET_LOCAL_ADDRESS:
                local_addr=(Address)evt.getArg();
                break;
            case Event.VIEW_CHANGE:
                handleViewChange(((View)evt.getArg()).getMembers());
                break;
            case Event.MSG:
                Message msg=(Message)evt.getArg();
                FcHeader hdr=(FcHeader)msg.removeHeader(name);
                if(hdr != null) {
                    switch(hdr.type) {
                    case FcHeader.REPLENISH:
                        num_replenishments++;
                        handleCredit(msg.getSrc());
                        break;
                    case FcHeader.CREDIT_REQUEST:
                        num_credit_requests++;
                        Address sender=msg.getSrc();
                        if(trace)
                            log.trace("received credit request from " + sender + ": sending credits");
                        received.put(sender, new Long(max_credits));
                        sendCredit(sender);
                        break;
                    default:
                        log.error("header type " + hdr.type + " not known");
                        break;
                    }
                    return; // don't pass message up
                }
                else {
                    adjustCredit(msg);
                }
                break;
        }
        passUp(evt);
    }



    private void handleCredit(Address sender) {
        if(sender == null) return;
        StringBuilder sb=null;
        boolean unblock=false;

        if(trace) {
            Long old_credit=(Long)sent.get(sender);
            sb=new StringBuilder();
            sb.append("received credit from ").append(sender).append(", old credit was ").
                    append(old_credit).append(", new credits are ").append(max_credits).
                    append(".\nCreditors before are: ").append(creditors);
        }

        synchronized(sent) {
            sent.put(sender, new Long(max_credits));
            if(creditors.size() > 0) {  // we are blocked because we expect credit from one or more members
                removeCreditor(sender);
                if(trace) {
                    sb.append("\nCreditors after removal of ").append(sender).append(" are: ").append(creditors);
                    log.trace(sb.toString());
                }
                if(creditors.size() == 0) {
                    unblock=true;
                }
            }
            else {  // no creditors, but still blocking: we need to unblock
                if(Boolean.TRUE.equals(blocking.get()))
                    unblock=true;
            }
        }
        if(unblock) // moved this outside of the 'sent' synchronized block
            unblockSender();
    }


    /**
     * Check whether sender has enough credits left. If not, send him some more
     * @param msg
     */
    private void adjustCredit(Message msg) {
        Address src=msg.getSrc();
        long    size=Math.max(24, msg.getLength());

        if(src == null) {
            if(log.isErrorEnabled()) log.error("src is null");
            return;
        }

        if(decrementCredit(received, src, size, min_credits) == false) {
            received.put(src, new Long(max_credits));
            if(trace) log.trace("sending replenishment message to " + src);
            sendCredit(src);
        }
    }



    private void sendCredit(Address dest) {
        Message  msg=new Message(dest, null, null);
        msg.putHeader(name, REPLENISH_HDR);
        passDown(new Event(Event.MSG, msg));
    }

    private void sendCreditRequest(final Address dest) {
        Message  msg=new Message(dest, null, null);
        msg.putHeader(name, CREDIT_REQUEST_HDR);
        passDown(new Event(Event.MSG, msg));
    }



    /**
     * Checks whether enough credits are available to send message. If not, blocks until enough credits
     * are available
     * @param evt Guaranteed to be a Message
     * @return
     */
    private void waitUntilEnoughCreditsAvailable() {
        while(true) {
            try {
                blocking.waitUntilWithTimeout(Boolean.FALSE, max_block_time);  // waits on 'sent'
                break;
            }
            catch(TimeoutException e) {
                List tmp=new ArrayList(creditors);
                if(trace)
                    log.trace("timeout occurred waiting for credits; sending credit request to " + tmp +
                              ", creditors are " + creditors);
                Address mbr;
                for(Iterator it=tmp.iterator(); it.hasNext();) {
                    mbr=(Address)it.next();
                    sendCreditRequest(mbr);
                }
            }
        }
    }


    /**
     * Try to decrement the credits needed for this message and return true if successful, or false otherwise.
     * For unicast destinations, the credits required are subtracted from the unicast destination member, for
     * multicast messages the credits are subtracted from all current members in the group.
     * @param msg
     * @return false: will block, true: will not block
     */
    private boolean decrMessage(Message msg) {
        Address dest;
        long    size;
        boolean success=true;

        // ******************************************************************************************************
        // this method is called by waitUntilEnoughCredits() which syncs on 'sent', so we don't need to sync here
        // ******************************************************************************************************

        if(msg == null) {
            if(log.isErrorEnabled()) log.error("msg is null");
            return true; // don't block !
        }
        dest=msg.getDest();
        size=Math.max(24, msg.getLength());
        if(dest != null && !dest.isMulticastAddress()) { // unicast destination
            if(decrementCredit(sent, dest, size, 0)) {
                return true;
            }
            else {
                addCreditor(dest);
                return false;
            }
        }
        else {                 // multicast destination
            for(Iterator it=members.iterator(); it.hasNext();) {
                dest=(Address)it.next();
                if(decrementCredit(sent, dest, size, 0) == false) {
                    addCreditor(dest);
                    success=false;
                }
            }
        }
        return success;
    }




    /** If message queueing is enabled, sends queued messages and unlocks sender (if successful) */
    private void unblockSender() {
        if(start_blocking > 0) {
            stop_blocking=System.currentTimeMillis();
            long diff=stop_blocking - start_blocking;
            total_time_blocking+=diff;
            last_blockings.add(new Long(diff));
            stop_blocking=start_blocking=0;
            if(trace)
                log.trace("setting blocking=false, blocking time was " + diff + "ms");
        }
        if(trace)
            log.trace("setting blocking=false");
        blocking.set(Boolean.FALSE);
    }


    private void addCreditor(Address mbr) {
        if(mbr != null && !creditors.contains(mbr))
            creditors.add(mbr);
    }

    private void removeCreditor(Address mbr) {
        creditors.remove(mbr);
    }




    /**
     * Find the credits associated with <tt>dest</tt> and decrement its credits by credits_required. If the remaining
     * value is less than or equal to 0, return false, else return true. Note that we will always subtract the credits.
     * @param map
     * @param dest
     * @param credits_required Number of bytes required
     * @param minimal_credits For the receiver: add minimal credits to check whether credits need to be sent
     * @return Whether the required credits could successfully be subtracted from the credits left
     */
    private boolean decrementCredit(Map map, Address dest, long credits_required, long minimal_credits) {
        long    credits_left, new_credits_left;
        Long    tmp=(Long)map.get(dest);
        boolean success;

        if(tmp == null)
            return true;

        credits_left=tmp.longValue();
        success=credits_left > (credits_required + minimal_credits);
        new_credits_left=Math.max(0, credits_left - credits_required);
        map.put(dest, new Long(new_credits_left));

        if(success) {
            return true;
        }
        else {
            if(trace) {
                StringBuilder sb=new StringBuilder();
                sb.append("not enough credits left for ").append(dest).append(": left=").append(new_credits_left);
                sb.append(", required+min_credits=").append((credits_required +min_credits)).append(", required=");
                sb.append(credits_required).append(", min_credits=").append(min_credits);
                log.trace(sb.toString());
            }
            return false;
        }
    }


    void handleViewChange(Vector mbrs) {
        Address addr;
        if(mbrs == null) return;

        if(trace) log.trace("new membership: " + mbrs);
        members.clear();
        members.addAll(mbrs);

        synchronized(received) {
            // add members not in membership to received hashmap (with full credits)
            for(int i=0; i < mbrs.size(); i++) {
                addr=(Address) mbrs.elementAt(i);
                if(!received.containsKey(addr))
                    received.put(addr, new Long(max_credits));
            }
            // remove members that left
            for(Iterator it=received.keySet().iterator(); it.hasNext();) {
                addr=(Address) it.next();
                if(!mbrs.contains(addr))
                    it.remove();
            }
        }

        boolean unblock=false;
        synchronized(sent) {
            // add members not in membership to sent hashmap (with full credits)
            for(int i=0; i < mbrs.size(); i++) {
                addr=(Address) mbrs.elementAt(i);
                if(!sent.containsKey(addr))
                    sent.put(addr, new Long(max_credits));
            }
            // remove members that left
            for(Iterator it=sent.keySet().iterator(); it.hasNext();) {
                addr=(Address)it.next();
                if(!mbrs.contains(addr))
                    it.remove(); // modified the underlying map
            }

            // remove all creditors which are not in the new view
            for(int i=0; i < creditors.size(); i++) {
                Address creditor=(Address)creditors.elementAt(i);
                if(!mbrs.contains(creditor))
                    creditors.remove(creditor);
            }

            if(trace) log.trace("creditors are " + creditors);
            if(creditors.size() == 0)
                unblock=true;
        }
        if(unblock)
            unblockSender();
    }

    private static String printMap(Map m) {
        Map.Entry entry;
        StringBuilder sb=new StringBuilder();
        for(Iterator it=m.entrySet().iterator(); it.hasNext();) {
            entry=(Map.Entry)it.next();
            sb.append(entry.getKey()).append(": ").append(entry.getValue()).append("\n");
        }
        return sb.toString();
    }





    public static class FcHeader extends Header implements Streamable {
        public static final byte REPLENISH      = 1;
        public static final byte CREDIT_REQUEST = 2; // the sender of the message is the requester

        byte  type = REPLENISH;

        public FcHeader() {

        }

        public FcHeader(byte type) {
            this.type=type;
        }

        public long size() {
            return Global.BYTE_SIZE;
        }

        public void writeExternal(ObjectOutput out) throws IOException {
            out.writeByte(type);
        }

        public void readExternal(ObjectInput in) throws IOException, ClassNotFoundException {
            type=in.readByte();
        }

        public void writeTo(DataOutputStream out) throws IOException {
            out.writeByte(type);
        }

        public void readFrom(DataInputStream in) throws IOException, IllegalAccessException, InstantiationException {
            type=in.readByte();
        }

        public String toString() {
            switch(type) {
            case REPLENISH: return "REPLENISH";
            case CREDIT_REQUEST: return "CREDIT_REQUEST";
            default: return "<invalid type>";
            }
        }
    }


}
