milvus/cmd/milvus/util.go
jaime c9d0c157ec
Move some modules from internal to public package (#22572)
Signed-off-by: jaime <yun.zhang@zilliz.com>
2023-04-06 19:14:32 +08:00

109 lines
2.5 KiB
Go

package milvus
import (
"fmt"
"io"
"io/ioutil"
"os"
"path"
"runtime"
"github.com/gofrs/flock"
"github.com/milvus-io/milvus/pkg/util/typeutil"
)
func makeRuntimeDir(dir string) error {
perm := os.FileMode(0755)
// os.MkdirAll equal to `mkdir -p`
err := os.MkdirAll(dir, perm)
if err != nil {
// err will be raised only when dir exists and dir is a file instead of a directory.
return fmt.Errorf("create runtime dir %s failed, err: %s", dir, err.Error())
}
tmpFile, err := ioutil.TempFile(dir, "tmp")
if err != nil {
return err
}
fileName := tmpFile.Name()
tmpFile.Close()
os.Remove(fileName)
return nil
}
// create runtime folder
func createRuntimeDir(sType string) string {
var writer io.Writer
if sType == typeutil.EmbeddedRole {
writer = io.Discard
} else {
writer = os.Stderr
}
runtimeDir := "/run/milvus"
if runtime.GOOS == "windows" {
runtimeDir = "run"
if err := makeRuntimeDir(runtimeDir); err != nil {
fmt.Fprintf(writer, "Create runtime directory at %s failed\n", runtimeDir)
os.Exit(-1)
}
} else {
if err := makeRuntimeDir(runtimeDir); err != nil {
fmt.Fprintf(writer, "Set runtime dir at %s failed, set it to /tmp/milvus directory\n", runtimeDir)
runtimeDir = "/tmp/milvus"
if err = makeRuntimeDir(runtimeDir); err != nil {
fmt.Fprintf(writer, "Create runtime directory at %s failed\n", runtimeDir)
os.Exit(-1)
}
}
}
return runtimeDir
}
func createPidFile(w io.Writer, filename string, runtimeDir string) (*flock.Flock, error) {
fileFullName := path.Join(runtimeDir, filename)
fd, err := os.OpenFile(fileFullName, os.O_CREATE|os.O_RDWR, 0664)
if err != nil {
return nil, fmt.Errorf("file %s is locked, error = %w", filename, err)
}
fmt.Fprintln(w, "open pid file:", fileFullName)
defer fd.Close()
fd.Truncate(0)
_, err = fd.WriteString(fmt.Sprintf("%d", os.Getpid()))
if err != nil {
return nil, fmt.Errorf("file %s write fail, error = %w", filename, err)
}
lock := flock.New(fileFullName)
_, err = lock.TryLock()
if err != nil {
return nil, fmt.Errorf("file %s is locked, error = %w", filename, err)
}
fmt.Fprintln(w, "lock pid file:", fileFullName)
return lock, nil
}
func getPidFileName(serverType string, alias string) string {
var filename string
if len(alias) != 0 {
filename = fmt.Sprintf("%s-%s.pid", serverType, alias)
} else {
filename = serverType + ".pid"
}
return filename
}
func closePidFile(fd *os.File) {
fd.Close()
}
func removePidFile(lock *flock.Flock) {
filename := lock.Path()
lock.Close()
os.Remove(filename)
}