package alluxio.worker.file;

import alluxio.AlluxioURI;
import alluxio.exception.ExceptionMessage;
import alluxio.exception.FileAlreadyExistsException;
import alluxio.exception.FileDoesNotExistException;
import alluxio.underfs.UnderFileSystem;
import alluxio.worker.WorkerContext;

import java.io.IOException;
import java.io.OutputStream;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

/**
 * Manages all open output streams to an Under File System. Individual streams should only be used
 * by one client.
 *
 * TODO(calvin): Enforce only one writer to each stream at a time
 */
public class UnderFileSystemManager {
  private static ConcurrentMap<AlluxioURI, UnderFileSystemManager> sUriToStreamManagerMap =
      new ConcurrentHashMap<>();

  private UnderFileSystem mUnderFileSystem;
  private ConcurrentMap<String, OutputStream> mFileToOutputStreamMap;

  public static UnderFileSystemManager get(AlluxioURI uri) {
    UnderFileSystemManager mgr = sUriToStreamManagerMap.get(uri);
    if (mgr == null) {
      UnderFileSystemManager newMgr = new UnderFileSystemManager(uri);
      mgr = sUriToStreamManagerMap.putIfAbsent(uri, newMgr);
      if (mgr == null) {
        mgr = newMgr;
      }
    }
    return mgr;
  }

  private UnderFileSystemManager(AlluxioURI uri) {
    mUnderFileSystem = UnderFileSystem.get(uri.toString(), WorkerContext.getConf());
    mFileToOutputStreamMap = new ConcurrentHashMap<>();
  }

  public OutputStream createFile(String path) throws FileAlreadyExistsException, IOException {
    OutputStream stream = mFileToOutputStreamMap.get(path);
    if (stream == null) {
      OutputStream newStream = mUnderFileSystem.create(path);
      stream = mFileToOutputStreamMap.putIfAbsent(path, newStream);
      if (stream == null) {
        return newStream;
      }
    }
    throw new FileAlreadyExistsException(ExceptionMessage.FAILED_UFS_CREATE.getMessage(path));
  }

  public void cancelFile(String path) throws FileDoesNotExistException, IOException {
    closeFile(path);
    mUnderFileSystem.delete(path, false);
  }

  public void closeFile(String path) throws FileDoesNotExistException, IOException {
    OutputStream stream = mFileToOutputStreamMap.remove(path);
    if (stream != null) {
      stream.close();
    } else {
      throw
          new FileDoesNotExistException(ExceptionMessage.UFS_PATH_DOES_NOT_EXIST.getMessage(path));
    }
  }
}

@@ -27,35 +27,8 @@
 e;%0A%0A
-import alluxio.AlluxioURI;%0A
 impo
@@ -403,16 +403,57 @@
 %0A%0A/**%0A *
+ Handles writes to the under file system.
  Manages
@@ -457,12 +457,8 @@
 ges 
-all 
 open
@@ -480,11 +480,8 @@
  to 
-an 
 Unde
@@ -493,17 +493,21 @@
 e System
-.
+s.%0A *
  Individ
@@ -537,19 +537,16 @@
  be used
-%0A *
  by one 
@@ -556,289 +556,118 @@
 ent.
-%0A *%0A * TODO(calvin): Enforce only one writer to each stream at a time%0A */%0Apublic class UnderFileSystemManager %7B%0A  private static ConcurrentMap%3CAlluxioURI, UnderFileSystemManager%3E sUriToStreamManagerMap =%0A      new ConcurrentHashMap%3C%3E();%0A%0A  private UnderFileSystem mUnderFileSystem;
+ Each instance keeps its own state of%0A * file names to open streams.%0A */%0Apublic class UnderFileSystemManager %7B
 %0A  p
@@ -747,15 +747,8 @@
 lic 
-static 
 Unde
@@ -769,469 +769,12 @@
 ager
- get(AlluxioURI uri) %7B%0A    UnderFileSystemManager mgr = sUriToStreamManagerMap.get(uri);%0A    if (mgr == null) %7B%0A      UnderFileSystemManager newMgr = new UnderFileSystemManager(uri);%0A      mgr = sUriToStreamManagerMap.putIfAbsent(uri, newMgr);%0A      if (mgr == null) %7B%0A        mgr = newMgr;%0A      %7D%0A    %7D%0A    return mgr;%0A  %7D%0A%0A  private UnderFileSystemManager(AlluxioURI uri) %7B%0A    mUnderFileSystem = UnderFileSystem.get(uri.toString(), WorkerContext.getConf());
+() %7B
 %0A   
@@ -1008,24 +1008,104 @@
  == null) %7B%0A
+      UnderFileSystem ufs = UnderFileSystem.get(path, WorkerContext.getConf());%0A
       Output
@@ -1123,32 +1123,19 @@
 tream = 
-mUnderFileSystem
+ufs
 .create(
@@ -1476,16 +1476,24 @@
 loseFile
+Internal
 (path);%0A
@@ -1500,74 +1500,267 @@
     
-mUnderFileSystem.delete(path, false);%0A  %7D%0A%0A  public void closeFile
+UnderFileSystem ufs = UnderFileSystem.get(path, WorkerContext.getConf());%0A    ufs.delete(path, false);%0A  %7D%0A%0A  public void closeFile(String path) throws FileDoesNotExistException, IOException %7B%0A    closeFileInternal(path);%0A  %7D%0A%0A  private void closeFileInternal
 (Str

Clean up under file system manager.