tools/stress/src/org/apache/cassandra/stress/generate/SeedManager.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.cassandra.stress.generate;

import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentSkipListMap;
import java.util.concurrent.atomic.AtomicLong;

import org.apache.cassandra.stress.Operation;
import org.apache.cassandra.stress.settings.StressSettings;
import org.apache.cassandra.utils.LockedDynamicList;

public class SeedManager
{

    final Distribution visits;
    final Generator writes;
    final Generator reads;
    final ConcurrentHashMap<Long, Seed> managing = new ConcurrentHashMap<>();
    final LockedDynamicList<Seed> sampleFrom;
    final Distribution sample;
    final long sampleOffset;
    final int sampleSize;
    final boolean updateSampleImmediately;

    public SeedManager(StressSettings settings)
    {
        Generator writes, reads;
        if (settings.generate.sequence != null)
        {
            long[] seq = settings.generate.sequence;
            if (settings.generate.readlookback != null)
            {
                LookbackableWriteGenerator series = new LookbackableWriteGenerator(seq[0], seq[1], settings.generate.wrap, settings.generate.readlookback.get());
                writes = series;
                reads = series.reads;
            }
            else
            {
                writes = reads = new SeriesGenerator(seq[0], seq[1], settings.generate.wrap);
            }
        }
        else
        {
            writes = reads = new RandomGenerator(settings.generate.distribution.get());
        }
        this.visits = settings.insert.visits.get();
        this.writes = writes;
        this.reads = reads;
        Distribution sample = settings.insert.revisit.get();
        this.sampleOffset = Math.min(sample.minValue(), sample.maxValue());
        long sampleSize = 1 + Math.max(sample.minValue(), sample.maxValue()) - sampleOffset;
        if (sampleOffset < 0 || sampleSize > Integer.MAX_VALUE)
            throw new IllegalArgumentException("sample range is invalid");
        this.sampleFrom = new LockedDynamicList<>((int) sampleSize);
        this.sample = DistributionInverted.invert(sample);
        this.sampleSize = (int) sampleSize;
        this.updateSampleImmediately = visits.average() > 1;
    }

    public Seed next(Operation op)
    {
        if (!op.isWrite())
        {
            Seed seed = reads.next(-1);
            if (seed == null)
                return null;
            Seed managing = this.managing.get(seed);
            return managing == null ? seed : managing;
        }

        while (true)
        {
            int index = (int) (sample.next() - sampleOffset);
            Seed seed = sampleFrom.get(index);
            if (seed != null && seed.isSaved())
                return seed;

            seed = writes.next((int) visits.next());
            if (seed == null)
                return null;
            if (managing.putIfAbsent(seed.seed, seed) == null)
            {
                if (!updateSampleImmediately || seed.save(sampleFrom, sampleSize))
                    return seed;
                managing.remove(seed.seed, seed);
            }
        }
    }

    public void markLastWrite(Seed seed, boolean first)
    {
        // we could have multiple iterators mark the last write simultaneously,
        // so we ensure we remove conditionally, and only remove the exact seed we were operating over
        // this is important because, to ensure correctness, we do not support calling remove multiple
        // times on the same DynamicList.Node
        if (managing.remove(seed.seed, seed) && !first)
            seed.remove(sampleFrom);
    }

    public void markFirstWrite(Seed seed, boolean last)
    {
        if (!last && !updateSampleImmediately)
            seed.save(sampleFrom, Integer.MAX_VALUE);
        writes.finishWrite(seed);
    }

    private abstract class Generator
    {
        abstract Seed next(int visits);
        void finishWrite(Seed seed) { }
    }

    private class RandomGenerator extends Generator
    {

        final Distribution distribution;

        public RandomGenerator(Distribution distribution)
        {
            this.distribution = distribution;
        }

        public Seed next(int visits)
        {
            return new Seed(distribution.next(), visits);
        }
    }

    private class SeriesGenerator extends Generator
    {

        final long start;
        final long totalCount;
        final boolean wrap;
        final AtomicLong next = new AtomicLong();

        public SeriesGenerator(long start, long end, boolean wrap)
        {
            this.wrap = wrap;
            if (start > end)
                throw new IllegalStateException();
            this.start = start;
            this.totalCount = 1 + end - start;
        }

        public Seed next(int visits)
        {
            long next = this.next.getAndIncrement();
            if (!wrap && next >= totalCount)
                return null;
            return new Seed(start + (next % totalCount), visits);
        }
    }

    private class LookbackableWriteGenerator extends SeriesGenerator
    {

        final AtomicLong writeCount = new AtomicLong();
        final ConcurrentSkipListMap<Seed, Seed> afterMin = new ConcurrentSkipListMap<>();
        final LookbackReadGenerator reads;

        public LookbackableWriteGenerator(long start, long end, boolean wrap, Distribution readLookback)
        {
            super(start, end, wrap);
            this.writeCount.set(0);
            reads = new LookbackReadGenerator(readLookback);
        }

        public Seed next(int visits)
        {
            long next = this.next.getAndIncrement();
            if (!wrap && next >= totalCount)
                return null;
            return new Seed(start + (next % totalCount), visits);
        }

        void finishWrite(Seed seed)
        {
            if (seed.seed <= writeCount.get())
                return;
            afterMin.put(seed, seed);
            while (true)
            {
                Map.Entry<Seed, Seed> head = afterMin.firstEntry();
                if (head == null)
                    return;
                long min = this.writeCount.get();
                if (head.getKey().seed <= min)
                    return;
                if (head.getKey().seed == min + 1 && this.writeCount.compareAndSet(min, min + 1))
                {
                    afterMin.remove(head.getKey());
                    continue;
                }
                return;
            }
        }

        private class LookbackReadGenerator extends Generator
        {

            final Distribution lookback;

            public LookbackReadGenerator(Distribution lookback)
            {
                this.lookback = lookback;
                if (lookback.maxValue() > start + totalCount)
                    throw new IllegalArgumentException("Invalid lookback distribution; max value is " + lookback.maxValue()
                                                       + ", but series only ranges from " + writeCount + " to " + (start + totalCount));
            }

            public Seed next(int visits)
            {
                long lookback = this.lookback.next();
                long range = writeCount.get();
                long startOffset = range - lookback;
                if (startOffset < 0)
                {
                    if (range == totalCount && !wrap)
                        return null;
                    startOffset = range == 0 ? 0 : lookback % range;
                }
                return new Seed(start + startOffset, visits);
            }
        }

    }

}