org.voltcore.zk.ZKCountdownLatch Maven / Gradle / Ivy

/* This file is part of VoltDB.
 * Copyright (C) 2008-2020 VoltDB Inc.
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program 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 Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with VoltDB.  If not, see .
 */

package org.voltcore.zk;

import org.apache.zookeeper_voltpatches.KeeperException;
import org.apache.zookeeper_voltpatches.ZooKeeper;
import org.apache.zookeeper_voltpatches.data.Stat;
import java.nio.ByteBuffer;

/**
 * A minimal countdown latch implementation that is not particularly scalable
 * nor good at cleaning up after itself. However it is fine for once per cluster type barriers
 * at startup although preferably it shouldn't be used in the common case where it would slow down unit
 * test startup.
 *
 * Initially created for a barrier after command logs are initialized and CL isn't on in most unit tests
 */
public class ZKCountdownLatch
{
    private final ZooKeeper m_zk;
    private final String m_path;
    private boolean countedDown = false;

    public ZKCountdownLatch(ZooKeeper zk, String path, int count) throws InterruptedException, KeeperException {
        m_zk = zk;
        m_path = path;

        ByteBuffer buf = ByteBuffer.allocate(4);
        buf.putInt(count);
        ZKUtil.asyncMkdirs(zk, path, buf.array());
    }

    public void await() throws InterruptedException, KeeperException {
        if (countedDown) return;
        while (true) {
            ZKUtil.FutureWatcher fw = new ZKUtil.FutureWatcher();
            int count = ByteBuffer.wrap(m_zk.getData(m_path, fw, null)).getInt();
            if (count > 0) {
                fw.get();
                continue;
            }
            countedDown = true;
            return;
        }
    }

    // Returns the current count
    public int getCount() throws InterruptedException, KeeperException {
        return ByteBuffer.wrap(m_zk.getData(m_path, false, null)).getInt();
    }

    // Returns if already counted down to zero
    public boolean isCountedDown() throws InterruptedException, KeeperException {
        if (countedDown) return true;
        int count = ByteBuffer.wrap(m_zk.getData(m_path, false, null)).getInt();
        if (count > 0) return false;
        countedDown = true;
        return true;
    }

    public void countDown() throws InterruptedException, KeeperException {
        countDown(false);
    }

    public void countDown(boolean expectNonZeroCount) throws InterruptedException, KeeperException {
        while (true) {
            Stat stat = new Stat();
            ByteBuffer buf = ByteBuffer.wrap(m_zk.getData(m_path, false, stat));
            int count = buf.getInt();
            if (count == 0) {
                countedDown = true;
                if (expectNonZeroCount) {
                    throw new RuntimeException("Count should be > 0");
                }
                return;
            }

            count--;

            //Save a few milliseconds
            if (count == 0) countedDown = true;

            buf.clear();
            buf.putInt(count);
            try {
                m_zk.setData(m_path, buf.array(), stat.getVersion());
            } catch (KeeperException.BadVersionException e) {
                continue;
            }
            return;
        }
    }
}