<commit_before>/*
 * Copyright (c) 2008-2016, Hazelcast, Inc. All Rights Reserved.
 *
 * Licensed 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 com.hazelcast.jet2.impl;

import com.hazelcast.jet2.JetEngineConfig;
import com.hazelcast.util.concurrent.BackoffIdleStrategy;
import com.hazelcast.util.concurrent.IdleStrategy;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.locks.LockSupport;

import static java.util.Collections.singletonList;
import static java.util.concurrent.TimeUnit.MICROSECONDS;
import static java.util.concurrent.TimeUnit.MILLISECONDS;
import static java.util.stream.Collectors.partitioningBy;

class ExecutionService {

    private static final IdleStrategy IDLER =
            new BackoffIdleStrategy(0, 0, MICROSECONDS.toNanos(1), MILLISECONDS.toNanos(1));
    private final ExecutorService blockingTaskletExecutor = Executors.newCachedThreadPool();
    private final Worker[] workers;
    private final Thread[] threads;

    public ExecutionService(JetEngineConfig cfg) {
        this.workers = new Worker[cfg.getParallelism()];
        this.threads = new Thread[cfg.getParallelism()];
    }

    public Future<Void> execute(List<Tasklet> tasklets) {
        final CountDownLatch completionLatch = new CountDownLatch(tasklets.size());
        final JobFuture jobFuture = new JobFuture(completionLatch);
        final Map<Boolean, List<Tasklet>> byBlocking = tasklets.stream().collect(partitioningBy(Tasklet::isBlocking));
        submitBlockingTasklets(completionLatch, jobFuture, byBlocking.get(true));
        submitNonblockingTasklets(completionLatch, jobFuture, byBlocking.get(false));
        return jobFuture;
    }

    private void submitBlockingTasklets(CountDownLatch completionLatch, JobFuture jobFuture, List<Tasklet> tasklets) {
        for (Tasklet t : tasklets) {
            blockingTaskletExecutor.execute(() -> {
                if (initPropagatingFailure(t, jobFuture, completionLatch)) {
                    blockingTaskletExecutor.execute(new Worker(new TaskletTracker(t, completionLatch, jobFuture)));
                }
            });
        }
    }

    private void submitNonblockingTasklets(CountDownLatch completionLatch, JobFuture jobFuture, List<Tasklet> tasklets) {
        ensureThreadsStarted();
        final List<TaskletTracker>[] trackersByThread = new List[workers.length];
        Arrays.setAll(trackersByThread, i -> new ArrayList());
        int i = 0;
        for (Tasklet t : tasklets) {
            if (initPropagatingFailure(t, jobFuture, completionLatch)) {
                trackersByThread[i++ % trackersByThread.length].add(new TaskletTracker(t, completionLatch, jobFuture));
            }
        }
        for (i = 0; i < trackersByThread.length; i++) {
            workers[i].trackers.addAll(trackersByThread[i]);
        }
        Arrays.stream(threads).forEach(LockSupport::unpark);
    }

    private boolean initPropagatingFailure(Tasklet t, JobFuture jobFuture, CountDownLatch completionLatch) {
        try {
            t.init();
            return true;
        } catch (Throwable e) {
            jobFuture.setTrouble(e);
            completionLatch.countDown();
            return false;
        }
    }

    private void ensureThreadsStarted() {
        if (workers[0] != null) {
            return;
        }
        Arrays.setAll(workers, i -> new Worker(workers));
        Arrays.setAll(threads, i -> new Thread(workers[i]));
        Arrays.stream(threads).forEach(Thread::start);
    }

    private static class Worker implements Runnable {
        private static final Worker[] EMPTY_WORKERS = new Worker[0];
        private final List<TaskletTracker> trackers;
        private final Worker[] colleagues;

        Worker(Worker[] colleagues) {
            this.colleagues = colleagues;
            this.trackers = new CopyOnWriteArrayList<>();
        }

        Worker(TaskletTracker tracker) {
            this.colleagues = EMPTY_WORKERS;
            this.trackers = singletonList(tracker);
        }

        @Override
        public void run() {
            long idleCount = 0;
            while (true) {
                boolean madeProgress = false;
                for (TaskletTracker t : trackers) {
                    final Worker stealingWorker = t.stealingWorker.get();
                    if (stealingWorker != null) {
                        t.stealingWorker.set(null);
                        trackers.remove(t);
                        stealingWorker.trackers.add(t);
                        continue;
                    }
                    try {
                        final ProgressState result = t.tasklet.call();
                        if (result.isDone()) {
                            dismissTasklet(t);
                        } else {
                            madeProgress |= result.isMadeProgress();
                        }
                    } catch (Throwable e) {
                        t.troubleSetter.setTrouble(e);
                    }
                    if (t.troubleSetter.hasTrouble()) {
                        dismissTasklet(t);
                    }
                }
                if (madeProgress) {
                    idleCount = 0;
                } else {
                    IDLER.idle(++idleCount);
                }
            }
        }

        private void dismissTasklet(TaskletTracker t) {
            t.completionLatch.countDown();
            trackers.remove(t);
            stealWork();
        }

        private void stealWork() {
            while (true) {
                // start with own tasklet list, try to find a longer one
                List<TaskletTracker> toStealFrom = trackers;
                for (Worker w : colleagues) {
                    if (w.trackers.size() > toStealFrom.size()) {
                        toStealFrom = w.trackers;
                    }
                }
                // if we couldn't find a list longer by more than one, there's nothing to steal
                if (toStealFrom.size() <= trackers.size() + 1) {
                    return;
                }
                // now we must find a task on this list which isn't already scheduled for moving
                for (TaskletTracker t : toStealFrom) {
                    if (t.stealingWorker.compareAndSet(null, this)) {
                        return;
                    }
                }
            }
        }
    }

    private static final class TaskletTracker {
        final Tasklet tasklet;
        final TroubleSetter troubleSetter;
        final CountDownLatch completionLatch;
        final AtomicReference<Worker> stealingWorker = new AtomicReference<>();<commit_msg>Specialize Blocking Worker<commit_after>/*
 * Copyright (c) 2008-2016, Hazelcast, Inc. All Rights Reserved.
 *
 * Licensed 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 com.hazelcast.jet2.impl;

import com.hazelcast.jet2.JetEngineConfig;
import com.hazelcast.util.concurrent.BackoffIdleStrategy;
import com.hazelcast.util.concurrent.IdleStrategy;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.Map;
import java.util.concurrent.CopyOnWriteArrayList;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Executors;
import java.util.concurrent.Future;
import java.util.concurrent.atomic.AtomicReference;
import java.util.concurrent.locks.LockSupport;

import static java.util.Collections.singletonList;
import static java.util.concurrent.TimeUnit.MICROSECONDS;
import static java.util.concurrent.TimeUnit.MILLISECONDS;
import static java.util.stream.Collectors.partitioningBy;

class ExecutionService {

    private static final IdleStrategy IDLER =
            new BackoffIdleStrategy(0, 0, MICROSECONDS.toNanos(1), MILLISECONDS.toNanos(1));
    private final ExecutorService blockingTaskletExecutor = Executors.newCachedThreadPool();
    private final NonBlockingWorker[] workers;
    private final Thread[] threads;

    public ExecutionService(JetEngineConfig cfg) {
        this.workers = new NonBlockingWorker[cfg.getParallelism()];
        this.threads = new Thread[cfg.getParallelism()];
    }

    public Future<Void> execute(List<Tasklet> tasklets) {
        final CountDownLatch completionLatch = new CountDownLatch(tasklets.size());
        final JobFuture jobFuture = new JobFuture(completionLatch);
        final Map<Boolean, List<Tasklet>> byBlocking = tasklets.stream().collect(partitioningBy(Tasklet::isBlocking));
        submitBlockingTasklets(completionLatch, jobFuture, byBlocking.get(true));
        submitNonblockingTasklets(completionLatch, jobFuture, byBlocking.get(false));
        return jobFuture;
    }

    private void submitBlockingTasklets(CountDownLatch completionLatch, JobFuture jobFuture, List<Tasklet> tasklets) {
        for (Tasklet t : tasklets) {
            blockingTaskletExecutor.execute(new BlockingWorker(new TaskletTracker(t, completionLatch, jobFuture)));
        }
    }

    private void submitNonblockingTasklets(CountDownLatch completionLatch, JobFuture jobFuture, List<Tasklet> tasklets) {
        ensureThreadsStarted();
        final List<TaskletTracker>[] trackersByThread = new List[workers.length];
        Arrays.setAll(trackersByThread, i -> new ArrayList());
        int i = 0;
        for (Tasklet t : tasklets) {
            if (initPropagatingFailure(t, jobFuture, completionLatch)) {
                trackersByThread[i++ % trackersByThread.length].add(new TaskletTracker(t, completionLatch, jobFuture));
            }
        }
        for (i = 0; i < trackersByThread.length; i++) {
            workers[i].trackers.addAll(trackersByThread[i]);
        }
        Arrays.stream(threads).forEach(LockSupport::unpark);
    }

    private void ensureThreadsStarted() {
        if (workers[0] != null) {
            return;
        }
        Arrays.setAll(workers, i -> new NonBlockingWorker(workers));
        Arrays.setAll(threads, i -> new Thread(workers[i]));
        Arrays.stream(threads).forEach(Thread::start);
    }

    private static boolean initPropagatingFailure(Tasklet t, JobFuture jobFuture, CountDownLatch completionLatch) {
        try {
            t.init();
            return true;
        } catch (Throwable e) {
            jobFuture.setTrouble(e);
            completionLatch.countDown();
            return false;
        }
    }

    private static class BlockingWorker implements Runnable {
        private final TaskletTracker tracker;

        private BlockingWorker(TaskletTracker tracker) {
            this.tracker = tracker;
        }

        @Override
        public void run() {
            final Tasklet t = tracker.tasklet;
            try {
                t.init();
                while (!t.call().isDone()) {
                }
            } catch (Throwable e) {
                tracker.troubleSetter.setTrouble(e);
            } finally {
                tracker.completionLatch.countDown();
            }
        }
    }

    private static class NonBlockingWorker implements Runnable {
        private static final NonBlockingWorker[] EMPTY_WORKERS = new NonBlockingWorker[0];
        private final List<TaskletTracker> trackers;
        private final NonBlockingWorker[] colleagues;

        NonBlockingWorker(NonBlockingWorker[] colleagues) {
            this.colleagues = colleagues;
            this.trackers = new CopyOnWriteArrayList<>();
        }

        NonBlockingWorker(TaskletTracker tracker) {
            this.colleagues = EMPTY_WORKERS;
            this.trackers = singletonList(tracker);
        }

        @Override
        public void run() {
            long idleCount = 0;
            while (true) {
                boolean madeProgress = false;
                for (TaskletTracker t : trackers) {
                    final NonBlockingWorker stealingWorker = t.stealingWorker.get();
                    if (stealingWorker != null) {
                        t.stealingWorker.set(null);
                        trackers.remove(t);
                        stealingWorker.trackers.add(t);
                        continue;
                    }
                    try {
                        final ProgressState result = t.tasklet.call();
                        if (result.isDone()) {
                            dismissTasklet(t);
                        } else {
                            madeProgress |= result.isMadeProgress();
                        }
                    } catch (Throwable e) {
                        t.troubleSetter.setTrouble(e);
                    }
                    if (t.troubleSetter.hasTrouble()) {
                        dismissTasklet(t);
                    }
                }
                if (madeProgress) {
                    idleCount = 0;
                } else {
                    IDLER.idle(++idleCount);
                }
            }
        }

        private void dismissTasklet(TaskletTracker t) {
            t.completionLatch.countDown();
            trackers.remove(t);
            stealWork();
        }

        private void stealWork() {
            while (true) {
                // start with own tasklet list, try to find a longer one
                List<TaskletTracker> toStealFrom = trackers;
                for (NonBlockingWorker w : colleagues) {
                    if (w.trackers.size() > toStealFrom.size()) {
                        toStealFrom = w.trackers;
                    }
                }
                // if we couldn't find a list longer by more than one, there's nothing to steal
                if (toStealFrom.size() <= trackers.size() + 1) {
                    return;
                }
                // now we must find a task on this list which isn't already scheduled for moving
                for (TaskletTracker t : toStealFrom) {
                    if (t.stealingWorker.compareAndSet(null, this)) {
                        return;
                    }
                }
            }
        }
    }

    private static final class TaskletTracker {
        final Tasklet tasklet;
        final TroubleSetter troubleSetter;
        final CountDownLatch completionLatch;
        final AtomicReference<NonBlockingWorker> stealingWorker = new AtomicReference<>();
{'fraction_non_alphanumeric': 0.08807327696643608, 'fraction_numerical': 0.003202664616961312, 'mean_word_length': 2.0322392697611185, 'pattern_counts': {'":': 0, '<': 31, '<?xml version=': 0, '>': 36, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3439454', 'n_tokens_mistral': 4204, 'n_tokens_neox': 3994, 'n_words': 972}