com.mapr.storm.SpoutState.java Source code

Introduction

Here is the source code for com.mapr.storm.SpoutState.java

Source

/*
 * Copyright MapR Technologies, $year
 *
 * 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.mapr.storm;

import com.google.common.base.Charsets;
import com.google.common.base.Preconditions;
import com.google.common.collect.Maps;
import com.google.common.io.Files;
import com.google.gson.Gson;
import com.google.gson.GsonBuilder;
import com.google.gson.JsonDeserializationContext;
import com.google.gson.JsonDeserializer;
import com.google.gson.JsonElement;
import com.google.gson.JsonObject;
import com.google.gson.JsonParseException;
import com.google.gson.reflect.TypeToken;
import com.mapr.storm.streamparser.StreamParser;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.File;
import java.io.IOException;
import java.lang.reflect.Type;
import java.util.Map;
import java.util.Queue;
import java.util.Random;
import java.util.Set;
import java.util.regex.Pattern;

/**
* Store and restore the spout state.
*/
public class SpoutState {
    private static final Logger log = LoggerFactory.getLogger(SpoutState.class);

    private Set<File> oldFiles;
    private String inputDirectory;
    private String filePattern;
    private Map<File, Long> offsets;

    public SpoutState() {
    }

    public Pattern getFilePattern() {
        return Pattern.compile(filePattern);
    }

    public String getInputDirectory() {
        return inputDirectory;
    }

    public Map<File, Long> getOffsets() {
        return offsets;
    }

    public Set<File> getOldFiles() {
        return oldFiles;
    }

    public static void recordCurrentState(Map<Long, PendingMessage> ackBuffer, DirectoryScanner scanner,
            StreamParser parser, File statusFile) {
        try {
            // find smallest offset for each file
            Map<File, Long> offsets = Maps.newHashMap();
            for (PendingMessage m : ackBuffer.values()) {
                Long x = offsets.get(m.getFile());
                if (x == null) {
                    x = m.getOffset();
                    offsets.put(m.getFile(), x);
                }
                if (m.getOffset() < x) {
                    offsets.put(m.getFile(), x);
                }
            }

            Long x = offsets.get(scanner.getLiveFile());
            if (x == null) {
                offsets.put(scanner.getLiveFile(), parser.currentOffset());
            }

            final String jsonState = new Gson().toJson(new SpoutState(scanner, offsets));
            final File newState = new File(statusFile.getParentFile(),
                    String.format("%s-%06x", statusFile.getName(), new Random().nextInt()));
            Files.write(jsonState, newState, Charsets.UTF_8);
            Files.move(newState, statusFile);
        } catch (IOException e) {
            log.error(String.format("Unable to write status to %s", statusFile), e);
        }
    }

    public static DirectoryScanner restoreState(Queue<PendingMessage> pendingReplays, File statusFile)
            throws IOException {
        Preconditions.checkState(statusFile.exists(), "Status file not found %s", statusFile);
        SpoutState s = SpoutState.fromString(Files.toString(statusFile, Charsets.UTF_8));
        DirectoryScanner scanner = new DirectoryScanner(new File(s.inputDirectory), Pattern.compile(s.filePattern));

        // we always reset all of the old replays.  Even in reliable = false cases,
        // there will be one of these entries for the live file.
        scanner.setOldFiles(s.oldFiles);
        for (File file : s.offsets.keySet()) {
            pendingReplays.add(new PendingMessage(file, s.offsets.get(file), null));
        }
        return scanner;
    }

    public static SpoutState fromString(String serialized) {
        GsonBuilder gson = new GsonBuilder();
        gson.registerTypeAdapter(SpoutState.class, new SpoutStateAdapter());
        return gson.create().fromJson(serialized, SpoutState.class);
    }

    private SpoutState(DirectoryScanner scanner, Map<File, Long> offsets) {
        this.oldFiles = scanner.getOldFiles();
        this.inputDirectory = scanner.getInputDirectory().toString();
        this.filePattern = scanner.getFileNamePattern().toString();
        this.offsets = offsets;
    }

    private static class SpoutStateAdapter implements JsonDeserializer<SpoutState> {

        @Override
        public SpoutState deserialize(JsonElement json, Type type, JsonDeserializationContext context)
                throws JsonParseException {
            SpoutState r = new SpoutState();
            final JsonObject object = json.getAsJsonObject();
            r.oldFiles = context.deserialize(object.get("oldFiles"), new TypeToken<Set<File>>() {
            }.getType());
            r.offsets = Maps.newHashMap();
            for (Map.Entry<String, JsonElement> entry : object.get("offsets").getAsJsonObject().entrySet()) {
                r.offsets.put(new File(entry.getKey()), entry.getValue().getAsLong());
            }
            r.inputDirectory = object.get("inputDirectory").getAsString();
            r.filePattern = object.get("filePattern").getAsString();
            return r;
        }
    }
}