Source code

001/*
002 * Licensed to the Apache Software Foundation (ASF) under one
003 * or more contributor license agreements.  See the NOTICE file
004 * distributed with this work for additional information
005 * regarding copyright ownership.  The ASF licenses this file
006 * to you under the Apache License, Version 2.0 (the
007 * "License"); you may not use this file except in compliance
008 * with the License.  You may obtain a copy of the License at
009 *
010 *   http://www.apache.org/licenses/LICENSE-2.0
011 *
012 * Unless required by applicable law or agreed to in writing,
013 * software distributed under the License is distributed on an
014 * "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
015 * KIND, either express or implied.  See the License for the
016 * specific language governing permissions and limitations
017 * under the License.
018 */
019package org.apache.reef.io.storage.local;
020
021import org.apache.reef.io.storage.ScratchSpace;
022
023import java.io.File;
024import java.io.IOException;
025import java.util.Set;
026import java.util.concurrent.ConcurrentSkipListSet;
027import java.util.logging.Logger;
028
029public class LocalScratchSpace implements ScratchSpace {
030
031  private static final Logger LOG = Logger.getLogger(LocalScratchSpace.class.getName());
032
033  private final String jobName;
034  private final String evaluatorName;
035  private final Set<File> tempFiles = new ConcurrentSkipListSet<>();
036  /**
037   * Zero denotes "unlimited".
038   */
039  private long quota;
040
041  public LocalScratchSpace(final String jobName, final String evaluatorName) {
042    this.jobName = jobName;
043    this.evaluatorName = evaluatorName;
044    this.quota = 0;
045  }
046
047  public LocalScratchSpace(final String jobName, final String evaluatorName, final long quota) {
048    this.jobName = jobName;
049    this.evaluatorName = evaluatorName;
050    this.quota = quota;
051  }
052
053  public File newFile() {
054    final File ret;
055    try {
056      ret = File.createTempFile("reef-" + jobName + "-" + evaluatorName,
057          "tmp");
058    } catch (final IOException e) {
059      throw new RuntimeException(e);
060    }
061    tempFiles.add(ret);
062    return ret;
063  }
064
065  @Override
066  public long availableSpace() {
067    return quota;
068  }
069
070  @Override
071  public long usedSpace() {
072    long ret = 0;
073    for (final File f : tempFiles) {
074      try {
075        ret += f.length();
076      } catch (final SecurityException e) {
077        LOG.info("Fail to get file info:" + f.getAbsolutePath());
078      }
079    }
080    return ret;
081  }
082
083  @Override
084  public void delete() {
085    for (final File f : tempFiles) {
086      try {
087        if (!f.delete()) {
088          f.deleteOnExit();
089        }
090      } catch (final SecurityException e) {
091        throw new RuntimeException("Fail to delete file:" + f.getAbsolutePath(), e);
092      }
093    }
094    tempFiles.clear();
095  }
096
097}