mirror of
https://gitee.com/milvus-io/milvus.git
synced 2026-02-02 01:06:41 +08:00
125 lines
2.9 KiB
Go
125 lines
2.9 KiB
Go
package milvus
|
|
|
|
import (
|
|
"flag"
|
|
"fmt"
|
|
"io"
|
|
"io/ioutil"
|
|
syslog "log"
|
|
"os"
|
|
"path"
|
|
"runtime"
|
|
|
|
"github.com/gofrs/flock"
|
|
|
|
"go.uber.org/automaxprocs/maxprocs"
|
|
|
|
"github.com/milvus-io/milvus/internal/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
|
|
}
|
|
|
|
// Initialize maxprocs
|
|
func initMaxprocs(serverType string, flags *flag.FlagSet) {
|
|
if serverType == typeutil.EmbeddedRole {
|
|
flags.SetOutput(io.Discard)
|
|
// Initialize maxprocs while discarding log.
|
|
maxprocs.Set(maxprocs.Logger(nil))
|
|
} else {
|
|
// Initialize maxprocs.
|
|
maxprocs.Set(maxprocs.Logger(syslog.Printf))
|
|
}
|
|
}
|
|
|
|
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)
|
|
}
|