package water.init;
import water.H2O;
import water.Iced;
import water.persist.Persist.PersistEntry;
import water.persist.PersistManager;
import water.util.FileUtils;
import water.util.Log;
import water.util.StringUtils;
import java.io.*;
import java.util.concurrent.atomic.AtomicLong;
import java.util.regex.Pattern;
public class NodePersistentStorage {
final String NPS_DIR;
final String NPS_SEPARATOR;
public static class NodePersistentStorageEntry extends Iced {
public String _category;
public String _name;
public long _size;
public long _timestamp_millis;
}
public NodePersistentStorage(String npsDir) {
if (npsDir == null) {
NPS_DIR = null;
NPS_SEPARATOR = null;
return;
}
if (H2O.getPM().isHdfsPath(npsDir)) {
NPS_SEPARATOR = "/";
}
else {
NPS_SEPARATOR = File.separator;
}
String s = npsDir.toString();
if (s.startsWith("file://")) {
s = s.substring(7);
}
else if (s.startsWith("file:")) {
s = s.substring(5);
}
NPS_DIR = s;
}
private void validateGeneral() {
if (NPS_DIR == null) {
throw new IllegalArgumentException("NodePersistentStorage directory not specified (try setting -flow_dir)");
}
}
private void validateCategoryName(String categoryName) {
if (categoryName == null) {
throw new IllegalArgumentException("NodePersistentStorage category not specified");
}
if (! Pattern.matches("[\\-a-zA-Z0-9]+", categoryName)) {
throw new IllegalArgumentException("NodePersistentStorage illegal category (" + categoryName + ")");
}
}
private void validateKeyName(String keyName) {
if (keyName == null) {
throw new IllegalArgumentException("NodePersistentStorage name not specified");
}
if (! Pattern.matches("[\\-a-zA-Z0-9_ \\(\\)]+", keyName)) {
throw new IllegalArgumentException("NodePersistentStorage illegal name (" + keyName + ")");
}
}
public boolean configured() {
return (NPS_DIR != null);
}
public boolean exists(String categoryName) {
validateGeneral();
validateCategoryName(categoryName);
String dirName = NPS_DIR + NPS_SEPARATOR + categoryName;
return H2O.getPM().exists(dirName);
}
public boolean exists(String categoryName, String keyName) {
validateGeneral();
validateCategoryName(categoryName);
validateKeyName(keyName);
String fileName = NPS_DIR + NPS_SEPARATOR + categoryName + NPS_SEPARATOR + keyName;
return H2O.getPM().exists(fileName);
}
public void put(String categoryName, String keyName, InputStream is) {
Log.info("NPS put content category(" + categoryName + ") keyName(" + keyName + ")");
// Error checking
validateGeneral();
validateCategoryName(categoryName);
validateKeyName(keyName);
// Create common directories
PersistManager pm = H2O.getPM();
if (! pm.exists(NPS_DIR)) {
boolean success = pm.mkdirs(NPS_DIR);
if (! success) {
throw new RuntimeException("Could not make NodePersistentStorage directory (" + NPS_DIR + ")");
}
}
if (! pm.exists(NPS_DIR)) {
throw new RuntimeException("NodePersistentStorage directory does not exist (" + NPS_DIR + ")");
}
String tmpd = NPS_DIR + NPS_SEPARATOR + "_tmp";
if (! pm.exists(tmpd)) {
boolean success = pm.mkdirs(tmpd);
if (! success) {
throw new RuntimeException("Could not make NodePersistentStorage category directory (" + tmpd + ")");
}
}
if (! pm.exists(tmpd)) {
throw new RuntimeException("NodePersistentStorage category directory does not exist (" + tmpd + ")");
}
// Create category directory
String d2 = NPS_DIR + NPS_SEPARATOR + categoryName;
if (! pm.exists(d2)) {
boolean success = pm.mkdirs(d2);
if (! success) {
throw new RuntimeException("Could not make NodePersistentStorage category directory (" + d2 + ")");
}
}
if (! pm.exists(d2)) {
throw new RuntimeException("NodePersistentStorage category directory does not exist (" + d2 + ")");
}
// Create tmp file
String tmpf = tmpd + NPS_SEPARATOR + keyName;
OutputStream os = null;
try {
os = pm.create(tmpf, true);
FileUtils.copyStream(is, os, 1024);
}
finally {
if (os != null) {
try {
os.close();
}
catch (Exception e) {
Log.err(e);
}
}
}
// Make final spot available if needed, and move tmp file to final spot.
boolean success;
String realf = d2 + NPS_SEPARATOR + keyName;
if (pm.exists(realf)) {
success = pm.delete(realf);
if (! success) {
throw new RuntimeException("NodePersistentStorage delete failed (" + realf + ")");
}
}
success = pm.rename(tmpf, realf);
if (! success) {
throw new RuntimeException("NodePersistentStorage rename failed (" + tmpf + " -> " + realf + ")");
}
if (! pm.exists(realf)) {
throw new RuntimeException("NodePersistentStorage file does not exist (" + realf + ")");
}
Log.info("Put succeeded");
}
public void put(String categoryName, String keyName, String value) {
validateGeneral();
validateCategoryName(categoryName);
validateKeyName(keyName);
InputStream is = new ByteArrayInputStream(StringUtils.bytesOf(value));
put(categoryName, keyName, is);
}
public NodePersistentStorageEntry[] list(String categoryName) {
validateGeneral();
validateCategoryName(categoryName);
String dirName = NPS_DIR + NPS_SEPARATOR + categoryName;
PersistEntry[] arr1 = H2O.getPM().list(dirName);
NodePersistentStorageEntry[] arr2 = new NodePersistentStorageEntry[arr1.length];
for (int i = 0; i < arr1.length; i++) {
arr2[i] = new NodePersistentStorageEntry();
arr2[i]._category = categoryName;
arr2[i]._name = arr1[i]._name;
arr2[i]._size = arr1[i]._size;
arr2[i]._timestamp_millis = arr1[i]._timestamp_millis;
}
return arr2;
}
public String get_as_string(String categoryName, String keyName) {
validateGeneral();
validateCategoryName(categoryName);
validateKeyName(keyName);
String fileName = NPS_DIR + NPS_SEPARATOR + categoryName + NPS_SEPARATOR + keyName;
InputStream is = H2O.getPM().open(fileName);
ByteArrayOutputStream baos = new ByteArrayOutputStream();
byte[] buf = new byte[4096];
try {
int n = is.read(buf, 0, buf.length);
while (true) {
if (baos.size() > (1024L * 1024L * 1024L)) {
throw new RuntimeException("File too big (" + fileName + ")");
}
if (n < 0) {
return baos.toString();
}
baos.write(buf, 0, n);
n = is.read(buf, 0, buf.length);
}
}
catch (Exception e) {
throw new RuntimeException(e);
}
}
public long get_length(String categoryName, String keyName) {
validateGeneral();
validateCategoryName(categoryName);
validateKeyName(keyName);
String fileName = NPS_DIR + NPS_SEPARATOR + categoryName + NPS_SEPARATOR + keyName;
if (! H2O.getPM().exists(fileName)) {
throw new IllegalArgumentException("File not found (" + fileName + ")");
}
return H2O.getPM().length(fileName);
}
public InputStream get(String categoryName, String keyName, AtomicLong length) {
validateGeneral();
validateCategoryName(categoryName);
validateKeyName(keyName);
String fileName = NPS_DIR + NPS_SEPARATOR + categoryName + NPS_SEPARATOR + keyName;
if (length != null) {
length.set(H2O.getPM().length(fileName));
}
return H2O.getPM().open(fileName);
}
public void delete(String categoryName, String keyName) {
validateGeneral();
validateCategoryName(categoryName);
validateKeyName(keyName);
String fileName = NPS_DIR + NPS_SEPARATOR + categoryName + NPS_SEPARATOR + keyName;
if (! H2O.getPM().exists(fileName)) {
return;
}
boolean success = H2O.getPM().delete(fileName);
if (! success) {
throw new RuntimeException("NodePersistentStorage delete failed (" + fileName + ")");
}
}
}