<commit_before>/*
 * Copyright 2017 <NAME>.
 *
 * 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 net.syberia.storm.rabbitmq;

import com.rabbitmq.client.Channel;
import com.rabbitmq.client.GetResponse;
import java.io.IOException;
import java.util.Map;
import java.util.concurrent.TimeoutException;
import org.apache.storm.spout.SpoutOutputCollector;
import org.apache.storm.task.TopologyContext;
import org.apache.storm.topology.OutputFieldsDeclarer;
import org.apache.storm.topology.base.BaseRichSpout;
import org.apache.storm.tuple.Fields;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 *
 * @author <NAME>
 */
public class RabbitMqSpout extends BaseRichSpout {

    private static final long serialVersionUID = 614091429512483100L;

    private static final Logger LOGGER = LoggerFactory.getLogger(RabbitMqSpout.class);
    
    public static final String KEY_QUEUE_NAME = "rabbitmq.queue_name";
    public static final String KEY_PREFETCH_COUNT = "rabbitmq.prefetch_count";
    public static final String KEY_REQUEUE_ON_FAIL = "rabbitmq.requeue_on_fail";

    private final RabbitMqChannelProvider rabbitMqChannelProvider;
    private final RabbitMqMessageScheme rabbitMqMessageScheme;

    private RabbitMqInitializer initializer;

    private String queueName;
    private int prefetchCount;
    private boolean requeueOnFail;
    private SpoutOutputCollector collector;

    private boolean active;

    public RabbitMqSpout(RabbitMqChannelProvider rabbitMqChannelProvider,
            RabbitMqMessageScheme rabbitMqMessageScheme) {
        this.rabbitMqChannelProvider = rabbitMqChannelProvider;
        this.rabbitMqMessageScheme = rabbitMqMessageScheme;
    }

    public void setInitializer(RabbitMqInitializer initializer) {
        this.initializer = initializer;
    }

    @Override
    public void open(Map conf, TopologyContext context, SpoutOutputCollector collector) {
        this.queueName = ConfigFetcher.fetchStringProperty(conf, KEY_QUEUE_NAME);
        this.prefetchCount = ConfigFetcher.fetchIntegerProperty(conf, KEY_PREFETCH_COUNT, 10);
        this.requeueOnFail = ConfigFetcher.fetchBooleanProperty(conf, KEY_REQUEUE_ON_FAIL, false);
        this.collector = collector;
        
        this.rabbitMqMessageScheme.prepare(conf, context);

        try {
            this.rabbitMqChannelProvider.prepare();
        } catch (IOException | TimeoutException ex) {
            throw new RuntimeException("Unable to prepare RabbitMQ channel provider", ex);
        }

        if (initializer != null) {
            Channel channel = getChannel();
            if (channel != null) {
                try {
                    initializer.initialize(channel);
                } catch (IOException ex) {
                    throw new RuntimeException("Unable to execute initialization", ex);
                } finally {
                    rabbitMqChannelProvider.returnChannel(channel);
                }
            }
        }
    }

    @Override
    public void nextTuple() {
        if (!active) {
            return;
        }

        Channel channel = getChannel();
        if (channel == null) {
            return;
        }

        try {
            for (int emitted = 0; emitted < prefetchCount; emitted++) {
                GetResponse response = channel.basicGet(queueName, false);
                if (response == null) { // no message retrieved
                    return;
                } else {
                    long messageId = response.getEnvelope().getDeliveryTag();
                    StreamedTuple streamedTuple;
                    try {
                        streamedTuple = rabbitMqMessageScheme.convertToStreamedTuple(response);
                    } catch (Exception ex) {
                        LOGGER.error("Unable to convert RabbitMQ message", ex);
                        collector.reportError(ex);
                        channel.basicReject(messageId, false);
                        continue;
                    }
                    if (streamedTuple == null) {
                        LOGGER.debug("Filtered message with id: {}", messageId);
                        channel.basicAck(messageId, false);
                    } else {
                        collector.emit(streamedTuple.getStreamId(), streamedTuple.getTuple(), messageId);
                    }
                }
            }
        } catch (IOException ex) {
            LOGGER.error("Unable to execute channel command", ex);
            collector.reportError(ex);
        } finally {
            rabbitMqChannelProvider.returnChannel(channel);
        }
    }

    @Override
    public void ack(Object msgId) {
        processMessageId(msgId, (Channel channel, long deliveryTag) -> {
            channel.basicAck(deliveryTag, false);
        });
    }

    @Override
    public void fail(Object msgId) {
        processMessageId(msgId, (Channel channel, long deliveryTag) -> {
            channel.basicReject(deliveryTag, requeueOnFail);
        });
    }

    private void processMessageId(Object msgId, ChannelAction channelAction) {
        Channel channel = getChannel();
        if (channel == null) {
            return;
        }
        long deliveryTag = (long) msgId;
        try {
            channelAction.execute(channel, deliveryTag);
        } catch (IOException ex) {
            LOGGER.error("Unable to process message id: " + deliveryTag, ex);
            collector.reportError(ex);
        } finally {
            rabbitMqChannelProvider.returnChannel(channel);
        }
    }

    private interface ChannelAction {

        void execute(Channel channel, long deliveryTag) throws IOException;
    }

    private Channel getChannel() {
        try {
            return rabbitMqChannelProvider.getChannel();
        } catch (Exception ex) {
            LOGGER.error("Unable to get RabbitMQ channel from the provider", ex);
            collector.reportError(ex);
            return null;
        }
    }

    @Override
    public void declareOutputFields(OutputFieldsDeclarer declarer) {
        Map<String, Fields> streamsOutputFields = rabbitMqMessageScheme.getStreamsOutputFields();
        for (String streamId : streamsOutputFields.keySet()) {
            declarer.declareStream(streamId, streamsOutputFields.get(streamId));
        }
    }

    @Override
    public void deactivate() {
        this.active = false;
    }

    @Override
    public void activate() {
        this.active = true;
    }

    @Override
    public void close() {
        try {
            this.rabbitMqChannelProvider.cleanup();
        } catch (Exception ex) {
            LOGGER.error("Unable to cleanup RabbitMQ provider", ex);
        }
        this.rabbitMqMessageScheme.cleanup();
    }

}<commit_msg>Use entrySet iterator instead of keySet in declareOutputFields method<commit_after>/*
 * Copyright 2017 <NAME>.
 *
 * 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 net.syberia.storm.rabbitmq;

import com.rabbitmq.client.Channel;
import com.rabbitmq.client.GetResponse;
import java.io.IOException;
import java.util.Map;
import java.util.Map.Entry;
import java.util.concurrent.TimeoutException;
import org.apache.storm.spout.SpoutOutputCollector;
import org.apache.storm.task.TopologyContext;
import org.apache.storm.topology.OutputFieldsDeclarer;
import org.apache.storm.topology.base.BaseRichSpout;
import org.apache.storm.tuple.Fields;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 *
 * @author <NAME>
 */
public class RabbitMqSpout extends BaseRichSpout {

    private static final long serialVersionUID = 614091429512483100L;

    private static final Logger LOGGER = LoggerFactory.getLogger(RabbitMqSpout.class);
    
    public static final String KEY_QUEUE_NAME = "rabbitmq.queue_name";
    public static final String KEY_PREFETCH_COUNT = "rabbitmq.prefetch_count";
    public static final String KEY_REQUEUE_ON_FAIL = "rabbitmq.requeue_on_fail";

    private final RabbitMqChannelProvider rabbitMqChannelProvider;
    private final RabbitMqMessageScheme rabbitMqMessageScheme;

    private RabbitMqInitializer initializer;

    private String queueName;
    private int prefetchCount;
    private boolean requeueOnFail;
    private SpoutOutputCollector collector;

    private boolean active;

    public RabbitMqSpout(RabbitMqChannelProvider rabbitMqChannelProvider,
            RabbitMqMessageScheme rabbitMqMessageScheme) {
        this.rabbitMqChannelProvider = rabbitMqChannelProvider;
        this.rabbitMqMessageScheme = rabbitMqMessageScheme;
    }

    public void setInitializer(RabbitMqInitializer initializer) {
        this.initializer = initializer;
    }

    @Override
    public void open(Map conf, TopologyContext context, SpoutOutputCollector collector) {
        this.queueName = ConfigFetcher.fetchStringProperty(conf, KEY_QUEUE_NAME);
        this.prefetchCount = ConfigFetcher.fetchIntegerProperty(conf, KEY_PREFETCH_COUNT, 10);
        this.requeueOnFail = ConfigFetcher.fetchBooleanProperty(conf, KEY_REQUEUE_ON_FAIL, false);
        this.collector = collector;
        
        this.rabbitMqMessageScheme.prepare(conf, context);

        try {
            this.rabbitMqChannelProvider.prepare();
        } catch (IOException | TimeoutException ex) {
            throw new RuntimeException("Unable to prepare RabbitMQ channel provider", ex);
        }

        if (initializer != null) {
            Channel channel = getChannel();
            if (channel != null) {
                try {
                    initializer.initialize(channel);
                } catch (IOException ex) {
                    throw new RuntimeException("Unable to execute initialization", ex);
                } finally {
                    rabbitMqChannelProvider.returnChannel(channel);
                }
            }
        }
    }

    @Override
    public void nextTuple() {
        if (!active) {
            return;
        }

        Channel channel = getChannel();
        if (channel == null) {
            return;
        }

        try {
            for (int emitted = 0; emitted < prefetchCount; emitted++) {
                GetResponse response = channel.basicGet(queueName, false);
                if (response == null) { // no message retrieved
                    return;
                } else {
                    long messageId = response.getEnvelope().getDeliveryTag();
                    StreamedTuple streamedTuple;
                    try {
                        streamedTuple = rabbitMqMessageScheme.convertToStreamedTuple(response);
                    } catch (Exception ex) {
                        LOGGER.error("Unable to convert RabbitMQ message", ex);
                        collector.reportError(ex);
                        channel.basicReject(messageId, false);
                        continue;
                    }
                    if (streamedTuple == null) {
                        LOGGER.debug("Filtered message with id: {}", messageId);
                        channel.basicAck(messageId, false);
                    } else {
                        collector.emit(streamedTuple.getStreamId(), streamedTuple.getTuple(), messageId);
                    }
                }
            }
        } catch (IOException ex) {
            LOGGER.error("Unable to execute channel command", ex);
            collector.reportError(ex);
        } finally {
            rabbitMqChannelProvider.returnChannel(channel);
        }
    }

    @Override
    public void ack(Object msgId) {
        processMessageId(msgId, (Channel channel, long deliveryTag) -> {
            channel.basicAck(deliveryTag, false);
        });
    }

    @Override
    public void fail(Object msgId) {
        processMessageId(msgId, (Channel channel, long deliveryTag) -> {
            channel.basicReject(deliveryTag, requeueOnFail);
        });
    }

    private void processMessageId(Object msgId, ChannelAction channelAction) {
        Channel channel = getChannel();
        if (channel == null) {
            return;
        }
        long deliveryTag = (long) msgId;
        try {
            channelAction.execute(channel, deliveryTag);
        } catch (IOException ex) {
            LOGGER.error("Unable to process message id: " + deliveryTag, ex);
            collector.reportError(ex);
        } finally {
            rabbitMqChannelProvider.returnChannel(channel);
        }
    }

    private interface ChannelAction {

        void execute(Channel channel, long deliveryTag) throws IOException;
    }

    private Channel getChannel() {
        try {
            return rabbitMqChannelProvider.getChannel();
        } catch (Exception ex) {
            LOGGER.error("Unable to get RabbitMQ channel from the provider", ex);
            collector.reportError(ex);
            return null;
        }
    }

    @Override
    public void declareOutputFields(OutputFieldsDeclarer declarer) {
        Map<String, Fields> streamsOutputFields = rabbitMqMessageScheme.getStreamsOutputFields();
        for (Entry<String, Fields> entry : streamsOutputFields.entrySet()) {
            declarer.declareStream(entry.getKey(), entry.getValue());
        }
    }

    @Override
    public void deactivate() {
        this.active = false;
    }

    @Override
    public void activate() {
        this.active = true;
    }

    @Override
    public void close() {
        try {
            this.rabbitMqChannelProvider.cleanup();
        } catch (Exception ex) {
            LOGGER.error("Unable to cleanup RabbitMQ provider", ex);
        }
        this.rabbitMqMessageScheme.cleanup();
    }

}
{'fraction_non_alphanumeric': 0.0823939331830293, 'fraction_numerical': 0.0042358406777345085, 'mean_word_length': 2.2674107142857145, 'pattern_counts': {'":': 0, '<': 12, '<?xml version=': 0, '>': 14, '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': '1789871', 'n_tokens_mistral': 3910, 'n_tokens_neox': 3684, 'n_words': 931}