<commit_before>/*
 * Copyright 2014-2015 Quantiply Corporation. 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.quantiply.samza.task;

import com.codahale.metrics.Meter;
import com.codahale.metrics.MetricRegistry;
import com.quantiply.samza.ConfigConst;
import com.quantiply.samza.MetricAdaptor;
import com.quantiply.samza.metrics.StreamMetricFactory;
import com.quantiply.samza.metrics.EventStreamMetrics;
import com.quantiply.samza.admin.KafkaAdmin;
import com.quantiply.samza.admin.TaskInfo;
import com.quantiply.samza.metrics.StreamMetricRegistry;
import org.apache.avro.Schema;
import org.apache.avro.generic.IndexedRecord;
import org.apache.avro.specific.SpecificRecord;
import org.apache.samza.config.Config;
import org.apache.samza.config.ConfigException;
import org.apache.samza.job.JobRunner;
import org.apache.samza.system.*;
import org.apache.samza.task.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.nio.file.Paths;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;

public abstract class BaseTask implements InitableTask, StreamTask, ClosableTask {
    protected TaskInfo taskInfo;
    protected Config config;
    protected Logger logger = LoggerFactory.getLogger(new Object(){}.getClass().getEnclosingClass());

    private final Map<String, StreamMsgHandler> handlerMap = new HashMap<>();
    private Optional<StreamMsgHandler> defaultHandler = Optional.empty();
    private MetricAdaptor metricAdaptor;
    private ErrorHandler errorHandler;

    @FunctionalInterface
    public interface Process {
        void apply(IncomingMessageEnvelope envelope, MessageCollector collector, TaskCoordinator coordinator) throws Exception;
    }

    @FunctionalInterface
    public interface ProcessWithMetrics<M> {
        void apply(IncomingMessageEnvelope envelope, MessageCollector collector, TaskCoordinator coordinator, M customMetrics) throws Exception;
    }

    private class StreamMetrics {
        public final Meter processed;
        public final Meter dropped;

        public StreamMetrics(StreamMetricRegistry registry) {
            processed = registry.meter("processed");
            dropped = registry.meter("dropped");
        }
    }

    private class StreamMsgHandler {
        private Optional<String> name;
        private Process processFunc;
        private StreamMetrics metrics;

        public StreamMsgHandler(Optional<String> name, Process processFunc, StreamMetrics metrics) {
            this.name = name;
            this.processFunc = processFunc;
            this.metrics = metrics;
        }

        public Optional<String> getName() {
            return name;
        }

        public void process(IncomingMessageEnvelope envelope, MessageCollector collector, TaskCoordinator coordinator) throws Exception {
            try {
                processFunc.apply(envelope, collector, coordinator);
                metrics.processed.mark();
            }
            catch (Exception e) {
                if (logger.isInfoEnabled()) {
                    //NOTE - logging at info level because these can be too numerous in PRD
                    logger.info("Error handling message", e);
                }
                errorHandler.handleError(envelope, e);
                //If control reaches here, the message was dropped
                metrics.dropped.mark();
            }
        }
    }

    @Override
    public void init(Config config, TaskContext context) throws Exception {
        this.config = config;
        taskInfo = new TaskInfo(config, context);
        metricAdaptor = new MetricAdaptor(new MetricRegistry(), context.getMetricsRegistry(), ConfigConst.METRICS_GROUP_NAME);
        errorHandler = new ErrorHandler(config, taskInfo);
        errorHandler.start();

        _init(config, context, metricAdaptor);

        validateHandlers(context);
    }

    private void validateHandlers(TaskContext context) {
        if (defaultHandler.isPresent()) {
            return;
        }
        for (SystemStreamPartition ssp: context.getSystemStreamPartitions()) {
            if (!handlerMap.containsKey(ssp.getStream())) {
                throw new ConfigException(
                        String.format("Missing handler for stream: %s. Call registerDefaultHandler() or registerHandler() in _init().", ssp.getStream())
                );
            }
        }
    }

    @Override
    public final void process(IncomingMessageEnvelope envelope, MessageCollector collector, TaskCoordinator coordinator) throws Exception {
        taskInfo.setMDC(envelope);

        //Dispatch
        StreamMsgHandler handler;
        String streamName = envelope.getSystemStreamPartition().getStream();
        if (handlerMap.containsKey(streamName)) {
            handler = handlerMap.get(streamName);
        }
        else if (defaultHandler.isPresent()) {
            handler = defaultHandler.get();
        }
        else {
            throw new ConfigException("No handler for input stream: " + streamName);
        }

        handler.process(envelope, collector, coordinator);

        taskInfo.clearMDC();
    }

    @Override
    public final void close() throws Exception {
        Exception error = null;
        if (errorHandler != null) {
            try {
                errorHandler.stop();
            }
            catch (Exception e) {
                logger.error("Error stopping error handler", e);
                error = e;
            }
        }

        try {
            _close();<commit_msg>Added debug message<commit_after>/*
 * Copyright 2014-2015 Quantiply Corporation. 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.quantiply.samza.task;

import com.codahale.metrics.Meter;
import com.codahale.metrics.MetricRegistry;
import com.quantiply.samza.ConfigConst;
import com.quantiply.samza.MetricAdaptor;
import com.quantiply.samza.metrics.StreamMetricFactory;
import com.quantiply.samza.metrics.EventStreamMetrics;
import com.quantiply.samza.admin.KafkaAdmin;
import com.quantiply.samza.admin.TaskInfo;
import com.quantiply.samza.metrics.StreamMetricRegistry;
import org.apache.avro.Schema;
import org.apache.avro.generic.IndexedRecord;
import org.apache.avro.specific.SpecificRecord;
import org.apache.samza.config.Config;
import org.apache.samza.config.ConfigException;
import org.apache.samza.job.JobRunner;
import org.apache.samza.system.*;
import org.apache.samza.task.*;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.nio.file.Paths;
import java.util.HashMap;
import java.util.Map;
import java.util.Optional;

public abstract class BaseTask implements InitableTask, StreamTask, ClosableTask {
    protected TaskInfo taskInfo;
    protected Config config;
    protected Logger logger = LoggerFactory.getLogger(new Object(){}.getClass().getEnclosingClass());

    private final Map<String, StreamMsgHandler> handlerMap = new HashMap<>();
    private Optional<StreamMsgHandler> defaultHandler = Optional.empty();
    private MetricAdaptor metricAdaptor;
    private ErrorHandler errorHandler;

    @FunctionalInterface
    public interface Process {
        void apply(IncomingMessageEnvelope envelope, MessageCollector collector, TaskCoordinator coordinator) throws Exception;
    }

    @FunctionalInterface
    public interface ProcessWithMetrics<M> {
        void apply(IncomingMessageEnvelope envelope, MessageCollector collector, TaskCoordinator coordinator, M customMetrics) throws Exception;
    }

    private class StreamMetrics {
        public final Meter processed;
        public final Meter dropped;

        public StreamMetrics(StreamMetricRegistry registry) {
            processed = registry.meter("processed");
            dropped = registry.meter("dropped");
        }
    }

    private class StreamMsgHandler {
        private Optional<String> name;
        private Process processFunc;
        private StreamMetrics metrics;

        public StreamMsgHandler(Optional<String> name, Process processFunc, StreamMetrics metrics) {
            this.name = name;
            this.processFunc = processFunc;
            this.metrics = metrics;
        }

        public Optional<String> getName() {
            return name;
        }

        public void process(IncomingMessageEnvelope envelope, MessageCollector collector, TaskCoordinator coordinator) throws Exception {
            try {
                processFunc.apply(envelope, collector, coordinator);
                metrics.processed.mark();
            }
            catch (Exception e) {
                if (logger.isInfoEnabled()) {
                    //NOTE - logging at info level because these can be too numerous in PRD
                    logger.info("Error handling message", e);
                }
                errorHandler.handleError(envelope, e);
                //If control reaches here, the message was dropped
                metrics.dropped.mark();
            }
        }
    }

    @Override
    public void init(Config config, TaskContext context) throws Exception {
        this.config = config;
        taskInfo = new TaskInfo(config, context);
        metricAdaptor = new MetricAdaptor(new MetricRegistry(), context.getMetricsRegistry(), ConfigConst.METRICS_GROUP_NAME);
        errorHandler = new ErrorHandler(config, taskInfo);
        errorHandler.start();

        _init(config, context, metricAdaptor);

        validateHandlers(context);
    }

    private void validateHandlers(TaskContext context) {
        if (defaultHandler.isPresent()) {
            return;
        }
        for (SystemStreamPartition ssp: context.getSystemStreamPartitions()) {
            if (!handlerMap.containsKey(ssp.getStream())) {
                throw new ConfigException(
                        String.format("Missing handler for stream: %s. Call registerDefaultHandler() or registerHandler() in _init().", ssp.getStream())
                );
            }
        }
    }

    @Override
    public final void process(IncomingMessageEnvelope envelope, MessageCollector collector, TaskCoordinator coordinator) throws Exception {
        taskInfo.setMDC(envelope);

        //Dispatch
        StreamMsgHandler handler;
        String streamName = envelope.getSystemStreamPartition().getStream();
        if (handlerMap.containsKey(streamName)) {
            handler = handlerMap.get(streamName);
        }
        else if (defaultHandler.isPresent()) {
            handler = defaultHandler.get();
        }
        else {
            throw new ConfigException("No handler for input stream: " + streamName);
        }

        handler.process(envelope, collector, coordinator);

        taskInfo.clearMDC();
    }

    @Override
    public final void close() throws Exception {
        logger.debug("Shutting down");

        Exception error = null;
        if (errorHandler != null) {
            try {
                errorHandler.stop();
            }
            catch (Exception e) {
                logger.error("Error stopping error handler", e);
                error = e;
            }
        }

        try {
            _close();
{'fraction_non_alphanumeric': 0.08435910060725423, 'fraction_numerical': 0.0022977186935828, 'mean_word_length': 2.7008806559368357, 'pattern_counts': {'":': 0, '<': 17, '<?xml version=': 0, '>': 17, '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': '2888787', 'n_tokens_mistral': 3111, 'n_tokens_neox': 2978, 'n_words': 843}