From fa072e373f6a191be5eeea0c4b57e8887f0282c6 Mon Sep 17 00:00:00 2001 From: criyle Date: Wed, 4 Mar 2020 16:41:52 -0500 Subject: [PATCH] add loacl storeage driver for executor server --- README.md | 3 + cmd/executorserver/file.go | 81 +++---------------------- cmd/executorserver/file_handler.go | 27 ++++++--- cmd/executorserver/file_local.go | 94 ++++++++++++++++++++++++++++++ cmd/executorserver/file_memory.go | 74 +++++++++++++++++++++++ cmd/executorserver/main.go | 11 ++++ cmd/executorserver/worker.go | 8 +-- 7 files changed, 214 insertions(+), 84 deletions(-) create mode 100644 cmd/executorserver/file_local.go create mode 100644 cmd/executorserver/file_memory.go diff --git a/README.md b/README.md index 7460d20..d23d8db 100644 --- a/README.md +++ b/README.md @@ -13,6 +13,7 @@ The goal to to reimplement [syzoj/judge-v3](https://github.com/syzoj/judge-v3) i A rest service to run program in restricted environment and it is basically a wrapper for `pkg/envexec` to run single / multiple programs. - /run POST execute program in the restricted environment +- /file GET list all cached file - /file POST prepare a file in the executor service (in memory), returns fileId (can be referenced in /run parameter) - /file/:fileId GET downloads file from executor service (in memory), returns file content - /file/:fileId DELETE delete file specified by fileId @@ -32,6 +33,8 @@ The default binding address for the executor server is `:5050`. Can be specified The default concurrency is `4`, Can be specified with `-parallism` flag. +The default file store is in memory, local cache can be specified wieh `-dir` flag. + #### Planed API interface ```typescript diff --git a/cmd/executorserver/file.go b/cmd/executorserver/file.go index 0c5348a..500518e 100644 --- a/cmd/executorserver/file.go +++ b/cmd/executorserver/file.go @@ -3,24 +3,18 @@ package main import ( "bytes" "crypto/rand" - "encoding/base64" - "sync" + "encoding/base32" + + "github.com/criyle/go-judge/file" ) const randIDLength = 12 -type fileData struct { - Content []byte - FileName string -} - -var ( - fileStore map[string]*fileData - fileStoreLock sync.RWMutex -) - -func init() { - fileStore = make(map[string]*fileData) +type fileStore interface { + Add(string, []byte) (string, error) // Add creates a file with name & content to the storage, returns id + Remove(string) bool // Remove deletes a file by id + Get(string) file.File // Get file by id, nil if not exists + List() []string // List return all file ids } func generateID() (string, error) { @@ -30,65 +24,8 @@ func generateID() (string, error) { } var buf bytes.Buffer - if _, err := base64.NewEncoder(base64.StdEncoding, &buf).Write(b); err != nil { + if _, err := base32.NewEncoder(base32.StdEncoding, &buf).Write(b); err != nil { return "", err } return string(buf.Bytes()), nil } - -func addFile(content []byte, fileName string) (string, error) { - fileStoreLock.Lock() - defer fileStoreLock.Unlock() - - var ( - id string - err error - ) - // generate until unique id (try maximun 50 times) - for i := 0; i < 50; i++ { - id, err = generateID() - if err != nil { - return "", err - } - if _, ok := fileStore[id]; !ok { - break - } - } - if err != nil { - return "", err - } - - fileStore[id] = &fileData{ - Content: content, - FileName: fileName, - } - return id, err -} - -func removeFile(fileID string) bool { - fileStoreLock.Lock() - defer fileStoreLock.Unlock() - - _, ok := fileStore[fileID] - delete(fileStore, fileID) - return ok -} - -func getFile(fileID string) (*fileData, bool) { - fileStoreLock.RLock() - defer fileStoreLock.RUnlock() - - f, ok := fileStore[fileID] - return f, ok -} - -func getAllFileID() []string { - fileStoreLock.RLock() - defer fileStoreLock.RUnlock() - - var b []string - for n := range fileStore { - b = append(b, n) - } - return b -} diff --git a/cmd/executorserver/file_handler.go b/cmd/executorserver/file_handler.go index 24aaa51..8c8bf66 100644 --- a/cmd/executorserver/file_handler.go +++ b/cmd/executorserver/file_handler.go @@ -1,6 +1,7 @@ package main import ( + "fmt" "io/ioutil" "mime" "net/http" @@ -10,7 +11,7 @@ import ( ) func fileGet(c *gin.Context) { - ids := getAllFileID() + ids := fs.List() c.JSON(http.StatusOK, ids) } @@ -32,8 +33,11 @@ func filePost(c *gin.Context) { return } - id, err := addFile(b, fh.Filename) - + id, err := fs.Add(fh.Filename, b) + if err != nil { + c.AbortWithError(http.StatusInternalServerError, err) + return + } c.JSON(http.StatusOK, id) } @@ -47,13 +51,20 @@ func fileIDGet(c *gin.Context) { return } - f, ok := getFile(uri.FileID) - if !ok { + f := fs.Get(uri.FileID) + if f == nil { c.AbortWithStatus(http.StatusNotFound) return } - typ := mime.TypeByExtension(path.Ext(f.FileName)) - c.Data(http.StatusOK, typ, f.Content) + content, err := f.Content() + if err != nil { + c.AbortWithError(http.StatusInternalServerError, err) + return + } + + typ := mime.TypeByExtension(path.Ext(f.Name())) + c.Header("Content-Disposition", fmt.Sprintf("attachment; filename=\"%s\"", f.Name())) + c.Data(http.StatusOK, typ, content) } func fileIDDelete(c *gin.Context) { @@ -66,7 +77,7 @@ func fileIDDelete(c *gin.Context) { return } - ok := removeFile(uri.FileID) + ok := fs.Remove(uri.FileID) if !ok { c.AbortWithStatus(http.StatusNotFound) return diff --git a/cmd/executorserver/file_local.go b/cmd/executorserver/file_local.go new file mode 100644 index 0000000..344f3f1 --- /dev/null +++ b/cmd/executorserver/file_local.go @@ -0,0 +1,94 @@ +package main + +import ( + "io/ioutil" + "os" + "path" + "sync" + + "github.com/criyle/go-judge/file" +) + +var _ fileStore = &fileLocalStore{} + +type fileLocalStore struct { + dir string // directory to store file + name map[string]string // id to name mapping if exists + mu sync.RWMutex +} + +func newFileLocalStore(dir string) *fileLocalStore { + return &fileLocalStore{ + dir: dir, + name: make(map[string]string), + } +} + +func (s *fileLocalStore) Add(name string, content []byte) (string, error) { + s.mu.Lock() + defer s.mu.Unlock() + + var ( + id string + err error + ) + // generate until unique id (try maximun 50 times) + for i := 0; i < 50; i++ { + id, err = generateID() + if err != nil { + return "", err + } + if _, err := os.Stat(path.Join(s.dir, id)); err == nil { + break + } + } + if err != nil { + return "", err + } + err = ioutil.WriteFile(path.Join(s.dir, id), content, 0644) + if err != nil { + return "", err + } + s.name[id] = name + return id, err +} + +func (s *fileLocalStore) Get(id string) file.File { + s.mu.RLock() + defer s.mu.RUnlock() + + p := path.Join(s.dir, id) + if _, err := os.Stat(p); os.IsNotExist(err) { + return nil + } + name, ok := s.name[id] + if !ok { + name = id + } + return file.NewLocalFile(name, p) +} + +func (s *fileLocalStore) Remove(id string) bool { + s.mu.Lock() + defer s.mu.Unlock() + + delete(s.name, id) + p := path.Join(s.dir, id) + if _, err := os.Stat(p); os.IsNotExist(err) { + return false + } + os.Remove(p) + return true +} + +func (s *fileLocalStore) List() []string { + var names []string + fi, err := ioutil.ReadDir(s.dir) + if err != nil { + return nil + } + for _, f := range fi { + names = append(names, f.Name()) + } + return names +} diff --git a/cmd/executorserver/file_memory.go b/cmd/executorserver/file_memory.go new file mode 100644 index 0000000..8ffa1bc --- /dev/null +++ b/cmd/executorserver/file_memory.go @@ -0,0 +1,74 @@ +package main + +import ( + "sync" + + "github.com/criyle/go-judge/file" +) + +var _ fileStore = &fileMemoryStore{} + +type fileMemoryStore struct { + store map[string]file.File + mu sync.RWMutex +} + +func newFileMemoryStore() *fileMemoryStore { + return &fileMemoryStore{ + store: make(map[string]file.File), + } +} + +func (s *fileMemoryStore) Add(fileName string, content []byte) (string, error) { + s.mu.Lock() + defer s.mu.Unlock() + + var ( + id string + err error + ) + // generate until unique id (try maximun 50 times) + for i := 0; i < 50; i++ { + id, err = generateID() + if err != nil { + return "", err + } + if _, ok := s.store[id]; !ok { + break + } + } + if err != nil { + return "", err + } + + s.store[id] = file.NewMemFile(fileName, content) + return id, err +} + +func (s *fileMemoryStore) Remove(fileID string) bool { + s.mu.Lock() + defer s.mu.Unlock() + + _, ok := s.store[fileID] + delete(s.store, fileID) + return ok +} + +func (s *fileMemoryStore) Get(fileID string) file.File { + s.mu.RLock() + defer s.mu.RUnlock() + + f := s.store[fileID] + return f +} + +func (s *fileMemoryStore) List() []string { + s.mu.RLock() + defer s.mu.RUnlock() + + var b []string + for n := range s.store { + b = append(b, n) + } + return b +} diff --git a/cmd/executorserver/main.go b/cmd/executorserver/main.go index b49866b..0498a99 100644 --- a/cmd/executorserver/main.go +++ b/cmd/executorserver/main.go @@ -6,6 +6,7 @@ import ( "flag" "io/ioutil" "log" + "os" "sync/atomic" "syscall" @@ -21,9 +22,12 @@ var ( addr = flag.String("http", ":5050", "specifies the http binding address") parallism = flag.Int("parallism", 4, "control the # of concurrency execution") tmpFsParam = flag.String("tmpfs", "size=8m,nr_inodes=4k", "tmpfs mount data") + dir = flag.String("dir", "", "specifies direcotry to store file upload / download (in memory by default)") envPool envexec.EnvironmentPool cgroupPool envexec.CgroupPool + + fs fileStore ) func init() { @@ -33,6 +37,13 @@ func init() { func main() { flag.Parse() + if *dir == "" { + fs = newFileMemoryStore() + } else { + os.MkdirAll(*dir, 0755) + fs = newFileLocalStore(*dir) + } + root, err := ioutil.TempDir("", "dm") if err != nil { panic(err) diff --git a/cmd/executorserver/worker.go b/cmd/executorserver/worker.go index 6fcf1a5..b3df4ee 100644 --- a/cmd/executorserver/worker.go +++ b/cmd/executorserver/worker.go @@ -129,7 +129,7 @@ func workDoSingle(rc cmd) (res response) { if copyOutSet[name] { res.Files[name] = string(b) } else { - id, err := addFile(b, name) + id, err := fs.Add(name, b) if err != nil { res.Status = status(envexec.StatusFileError) res.Error = err.Error() @@ -182,11 +182,11 @@ func prepareCmdFile(f *cmdFile) (interface{}, error) { case f.Content != nil: return file.NewMemFile("file", []byte(*f.Content)), nil case f.FileID != nil: - fd, ok := getFile(*f.FileID) - if !ok { + fd := fs.Get(*f.FileID) + if fd == nil { return nil, fmt.Errorf("file not exists for %v", *f.FileID) } - return file.NewMemFile(fd.FileName, fd.Content), nil + return fd, nil case f.Max != nil && f.Name != nil: return envexec.PipeCollector{Name: *f.Name, SizeLimit: *f.Max}, nil default: