/*
 * Copyright (c) 2017, Red Hat Inc. All rights reserved.
 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
 *
 * This code is free software; you can redistribute it and/or modify it
 * under the terms of the GNU General Public License version 2 only, as
 * published by the Free Software Foundation.  Oracle designates this
 * particular file as subject to the "Classpath" exception as provided
 * by Oracle in the LICENSE file that accompanied this code.
 *
 * This code is distributed in the hope that it will be useful, but WITHOUT
 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
 * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
 * version 2 for more details (a copy is included in the LICENSE file that
 * accompanied this code).
 *
 * You should have received a copy of the GNU General Public License version
 * 2 along with this work; if not, write to the Free Software Foundation,
 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
 *
 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
 * or visit www.oracle.com if you need additional information or have any
 * questions.
 */
package org.openjdk.gcbench.util.ratelimit;

import java.util.concurrent.atomic.AtomicLongFieldUpdater;

public class TokenBucket implements RateLimiter {

    static final int QUANTA_PER_SEC = 5;
    static final int MS_PER_QUANTUM = 1000 / QUANTA_PER_SEC;

    static final AtomicLongFieldUpdater<TokenBucket> STATE =
            AtomicLongFieldUpdater.newUpdater(TokenBucket.class, "state");

    private final int tokensPerQuantum;
    private final long timeBase;
    private volatile long state;

    public TokenBucket(int ratePerSec) {
        this.tokensPerQuantum = Math.max(1, ratePerSec / QUANTA_PER_SEC);
        this.timeBase = System.currentTimeMillis();
    }

    private static int timestamp(long l) {
        return (int)(l >> 32);
    }

    private static int tokens(long l) {
        return (int)(l & 0x7FFFFFFF);
    }

    private static long pack(int timestamp, int tokens) {
        return ((long)timestamp << 32) + tokens;
    }

    @Override
    public void limit() {
        while (true) {
            int quantId = (int) ((System.currentTimeMillis() - timeBase) / MS_PER_QUANTUM);

            long cur = STATE.get(this);
            int time = timestamp(cur);
            int tokens = tokens(cur);

            if (time == quantId && tokens != 0) {
                // current quantum has tokens, try to claim and exit
                if (STATE.compareAndSet(this, cur, pack(quantId, tokens - 1))) {
                    return; // success
                } else {
                    continue; // immediate respin
                }
            } else if (time <= quantId) {
                // current or past quantum is empty, try to install a new one, and respin
                STATE.compareAndSet(this, cur, pack(quantId + 1, tokensPerQuantum));
            }

            // no rush: wait before respinning
            try {
                Thread.sleep(1);
            } catch (InterruptedException e) {
                // ignore
            }
        }
    }

}