curator-master/curator-recipes/src/main/java/org/apache/curator/framework/recipes/atomic/DistributedAtomicInteger.java
/*
 * 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.curator.framework.recipes.atomic;

import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import java.nio.BufferOverflowException;
import java.nio.BufferUnderflowException;
import java.nio.ByteBuffer;
import org.apache.curator.RetryPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.recipes.locks.InterProcessMutex;

/**
 * <p>A counter that attempts atomic increments. It first tries uses optimistic locking. If that fails,
 * an optional {@link InterProcessMutex} is taken. For both optimistic and mutex, a retry policy is used to
 * retry the increment.</p>
 *
 * <p>The various increment methods return an {@link AtomicValue} object. You must <b>always</b> check
 * {@link AtomicValue#succeeded()}. None of the methods (other than get()) are guaranteed to succeed.</p>
 */
public class DistributedAtomicInteger implements DistributedAtomicNumber<Integer> {
    private final DistributedAtomicValue value;

    /**
     * Creates in optimistic mode only - i.e. the promotion to a mutex is not done
     *
     * @param client the client
     * @param counterPath path to hold the value
     * @param retryPolicy the retry policy to use
     */
    public DistributedAtomicInteger(CuratorFramework client, String counterPath, RetryPolicy retryPolicy) {
        this(client, counterPath, retryPolicy, null);
    }

    /**
     * Creates in mutex promotion mode. The optimistic lock will be tried first using
     * the given retry policy. If the increment does not succeed, a {@link InterProcessMutex} will be tried
     * with its own retry policy
     *
     * @param client the client
     * @param counterPath path to hold the value
     * @param retryPolicy the retry policy to use
     * @param promotedToLock the arguments for the mutex promotion
     */
    public DistributedAtomicInteger(
            CuratorFramework client, String counterPath, RetryPolicy retryPolicy, PromotedToLock promotedToLock) {
        value = new DistributedAtomicValue(client, counterPath, retryPolicy, promotedToLock);
    }

    @Override
    public AtomicValue<Integer> get() throws Exception {
        return new AtomicInteger(value.get());
    }

    @Override
    public void forceSet(Integer newValue) throws Exception {
        value.forceSet(valueToBytes(newValue));
    }

    @Override
    public AtomicValue<Integer> compareAndSet(Integer expectedValue, Integer newValue) throws Exception {
        return new AtomicInteger(value.compareAndSet(valueToBytes(expectedValue), valueToBytes(newValue)));
    }

    @Override
    public AtomicValue<Integer> trySet(Integer newValue) throws Exception {
        return new AtomicInteger(value.trySet(valueToBytes(newValue)));
    }

    @Override
    public boolean initialize(Integer initialize) throws Exception {
        return value.initialize(valueToBytes(initialize));
    }

    /**
     * Add 1 to the current value and return the new value information. Remember to always
     * check {@link AtomicValue#succeeded()}.
     *
     * @return value info
     * @throws Exception ZooKeeper errors
     */
    @Override
    public AtomicValue<Integer> increment() throws Exception {
        return worker(1);
    }

    /**
     * Subtract 1 from the current value and return the new value information. Remember to always
     * check {@link AtomicValue#succeeded()}.
     *
     * @return value info
     * @throws Exception ZooKeeper errors
     */
    @Override
    public AtomicValue<Integer> decrement() throws Exception {
        return worker(-1);
    }

    /**
     * Add delta to the current value and return the new value information. Remember to always
     * check {@link AtomicValue#succeeded()}.
     *
     * @param delta amount to add
     * @return value info
     * @throws Exception ZooKeeper errors
     */
    @Override
    public AtomicValue<Integer> add(Integer delta) throws Exception {
        return worker(delta);
    }

    /**
     * Subtract delta from the current value and return the new value information. Remember to always
     * check {@link AtomicValue#succeeded()}.
     *
     * @param delta amount to subtract
     * @return value info
     * @throws Exception ZooKeeper errors
     */
    @Override
    public AtomicValue<Integer> subtract(Integer delta) throws Exception {
        return worker(-1 * delta);
    }

    @VisibleForTesting
    byte[] valueToBytes(Integer newValue) {
        Preconditions.checkNotNull(newValue, "newValue cannot be null");

        byte[] newData = new byte[4];
        ByteBuffer wrapper = ByteBuffer.wrap(newData);
        wrapper.putInt(newValue);
        return newData;
    }

    @VisibleForTesting
    int bytesToValue(byte[] data) {
        if ((data == null) || (data.length == 0)) {
            return 0;
        }
        ByteBuffer wrapper = ByteBuffer.wrap(data);
        try {
            return wrapper.getInt();
        } catch (BufferUnderflowException e) {
            throw value.createCorruptionException(data);
        } catch (BufferOverflowException e) {
            throw value.createCorruptionException(data);
        }
    }

    private AtomicValue<Integer> worker(final Integer addAmount) throws Exception {
        Preconditions.checkNotNull(addAmount, "addAmount cannot be null");

        MakeValue makeValue = new MakeValue() {
            @Override
            public byte[] makeFrom(byte[] previous) {
                int previousValue = (previous != null) ? bytesToValue(previous) : 0;
                int newValue = previousValue + addAmount;
                return valueToBytes(newValue);
            }
        };

        AtomicValue<byte[]> result = value.trySet(makeValue);
        return new AtomicInteger(result);
    }

    private class AtomicInteger implements AtomicValue<Integer> {
        private AtomicValue<byte[]> bytes;

        private AtomicInteger(AtomicValue<byte[]> bytes) {
            this.bytes = bytes;
        }

        @Override
        public boolean succeeded() {
            return bytes.succeeded();
        }

        @Override
        public Integer preValue() {
            return bytesToValue(bytes.preValue());
        }

        @Override
        public Integer postValue() {
            return bytesToValue(bytes.postValue());
        }

        @Override
        public AtomicStats getStats() {
            return bytes.getStats();
        }
    }
}
