CheckpointManagerConfig.java

// Copyright 2016 Twitter. 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.twitter.heron.ckptmgr;

import java.io.File;
import java.time.Duration;
import java.util.HashMap;
import java.util.Map;

import com.twitter.heron.common.basics.ByteAmount;
import com.twitter.heron.common.basics.TypeUtils;
import com.twitter.heron.common.config.ConfigReader;

public final class CheckpointManagerConfig {

  private Map<String, Object> config = new HashMap<>();

  private CheckpointManagerConfig(Builder build) {
    this.config = new HashMap<>(build.keyValues);
  }

  public static Builder newBuilder(boolean loadDefaults) {
    return Builder.create(loadDefaults);
  }

  @SuppressWarnings("unchecked")
  public Map<String, Object> getStatefulStorageConfig() {
    Object statefulStorageConfigObject = get(CheckpointManagerConfigKey.STORAGE_CONFIG);

    if (statefulStorageConfigObject instanceof Map) {
      return (Map<String, Object>) statefulStorageConfigObject;
    } else {
      throw new IllegalArgumentException(
          String.format("configs for stateful storage needs to be map, but is: %s",
              statefulStorageConfigObject.getClass().getName()));
    }
  }

  public String getStorageClassname() {
    return getString(CheckpointManagerConfigKey.STORAGE_CLASSNAME);
  }

  public ByteAmount getWriteBatchSize() {
    return getByteAmount(CheckpointManagerConfigKey.WRITE_BATCH_SIZE);
  }

  public Duration getWriteBatchTime() {
    return getDuration(CheckpointManagerConfigKey.WRITE_BATCH_TIME);
  }

  public ByteAmount getReadBatchSize() {
    return getByteAmount(CheckpointManagerConfigKey.READ_BATCH_SIZE);
  }

  public Duration getReadBatchTime() {
    return getDuration(CheckpointManagerConfigKey.READ_BATCH_TIME);
  }

  public ByteAmount getSocketSendSize() {
    return getByteAmount(CheckpointManagerConfigKey.SOCKET_SEND_SIZE);
  }

  public ByteAmount getSocketReceiveSize() {
    return getByteAmount(CheckpointManagerConfigKey.SOCKET_RECEIVE_SIZE);
  }

  private String getString(CheckpointManagerConfigKey key) {
    assertType(key, CheckpointManagerConfigKey.Type.STRING);
    return (String) get(key);
  }

  private Integer getInteger(CheckpointManagerConfigKey key) {
    assertType(key, CheckpointManagerConfigKey.Type.INTEGER);
    return TypeUtils.getInteger(get(key));
  }

  private Long getLong(CheckpointManagerConfigKey key) {
    assertType(key, CheckpointManagerConfigKey.Type.LONG);
    return TypeUtils.getLong(get(key));
  }

  private Duration getDuration(CheckpointManagerConfigKey key) {
    assertType(key, CheckpointManagerConfigKey.Type.DURATION);
    return TypeUtils.getDuration(get(key), key.getTemporalUnit());
  }

  private ByteAmount getByteAmount(CheckpointManagerConfigKey key) {
    assertType(key, CheckpointManagerConfigKey.Type.BYTE_AMOUNT);
    return TypeUtils.getByteAmount(get(key));
  }

  private Object get(CheckpointManagerConfigKey key) {
    return config.get(key.value());
  }

  private void assertType(CheckpointManagerConfigKey key, CheckpointManagerConfigKey.Type type) {
    if (key.getType() != type) {
      throw new IllegalArgumentException(String.format(
          "config key %s is of type %s instead of expected type %s", key, key.getType(), type));
    }
  }

  public static class Builder {
    private final Map<String, Object> keyValues = new HashMap<>();

    private static CheckpointManagerConfig.Builder create(boolean loadDefaults) {
      CheckpointManagerConfig.Builder cb = new Builder();

      if (loadDefaults) {
        loadDefaults(cb, CheckpointManagerConfigKey.values());
      }

      return cb;
    }

    private static void loadDefaults(CheckpointManagerConfig.Builder cb,
                                     CheckpointManagerConfigKey... keys) {
      for (CheckpointManagerConfigKey key : keys) {
        if (key.getDefault() != null) {
          cb.put(key, key.getDefault());
        }
      }
    }

    public Builder put(CheckpointManagerConfigKey key, Object value) {
      convertAndAdd(this.keyValues, key, value);
      return this;
    }

    public Builder putAll(String fileName, boolean mustExist) {
      File file = new File(fileName);
      if (!file.exists() && mustExist) {
        throw new IllegalArgumentException(
            String.format("Config file %s does not exist", fileName));
      }

      Map<String, Object> configValues = ConfigReader.loadFile(fileName);
      for (String keyValue : configValues.keySet()) {
        CheckpointManagerConfigKey key =
            CheckpointManagerConfigKey.toCheckpointManagerConfigKey(keyValue);
        if (key != null) {
          convertAndAdd(configValues, key, configValues.get(keyValue));
        }
      }
      keyValues.putAll(configValues);
      return this;
    }

    private static void convertAndAdd(Map<String, Object> config,
                                      CheckpointManagerConfigKey key,
                                      Object value) {
      if (key != null) {
        switch (key.getType()) {
          case BYTE_AMOUNT:
            config.put(key.value(), TypeUtils.getByteAmount(value));
            break;
          case INTEGER:
            config.put(key.value(), TypeUtils.getInteger(value));
            break;
          case LONG:
            config.put(key.value(), TypeUtils.getLong(value));
            break;
          case STRING:
            config.put(key.value(), value);
            break;
          case MAP:
            config.put(key.value(), value);
            break;
          default:
            throw new IllegalArgumentException(String.format(
                "config key %s is of type %s which is not yet supported", key, key.getType()));
        }
      }
    }

    public CheckpointManagerConfig build() {
      return new CheckpointManagerConfig(this);
    }
  }
}