dragonfly/cmd/scheduler/cmd/root.go

140 lines
3.6 KiB
Go

/*
* Copyright 2020 The Dragonfly Authors
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package cmd
import (
"context"
"fmt"
"os"
"path"
"github.com/spf13/cobra"
"d7y.io/dragonfly/v2/cmd/dependency"
logger "d7y.io/dragonfly/v2/internal/dflog"
"d7y.io/dragonfly/v2/pkg/dfpath"
"d7y.io/dragonfly/v2/pkg/types"
"d7y.io/dragonfly/v2/scheduler"
"d7y.io/dragonfly/v2/scheduler/config"
"d7y.io/dragonfly/v2/version"
)
var (
cfg *config.Config
)
// rootCmd represents the commonv1 command when called without any subcommands.
var rootCmd = &cobra.Command{
Use: "scheduler",
Short: "the scheduler of dragonfly",
Long: `Scheduler is a long-running process which receives and manages download tasks from the dfdaemon, notify the seed peer to return to the source,
generate and maintain a P2P network during the download process, and push suitable download nodes to the dfdaemon`,
Args: cobra.NoArgs,
DisableAutoGenTag: true,
SilenceUsage: true,
RunE: func(cmd *cobra.Command, args []string) error {
// Convert config.
if err := cfg.Convert(); err != nil {
return err
}
// Validate config.
if err := cfg.Validate(); err != nil {
return err
}
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
// Initialize dfpath.
d, err := initDfpath(&cfg.Server)
if err != nil {
return err
}
rotateConfig := logger.LogRotateConfig{
MaxSize: cfg.Server.LogMaxSize,
MaxAge: cfg.Server.LogMaxAge,
MaxBackups: cfg.Server.LogMaxBackups}
// Initialize logger.
if err := logger.InitScheduler(cfg.Verbose, cfg.Console, d.LogDir(), rotateConfig); err != nil {
return fmt.Errorf("init scheduler logger: %w", err)
}
logger.RedirectStdoutAndStderr(cfg.Console, path.Join(d.LogDir(), types.SchedulerName))
return runScheduler(ctx, d)
},
}
// Execute adds all child commands to the root command and sets flags appropriately.
// This is called by main.main(). It only needs to happen once to the rootCmd.
func Execute() {
if err := rootCmd.Execute(); err != nil {
logger.Error(err)
os.Exit(1)
}
}
func init() {
// Initialize default scheduler config.
cfg = config.New()
// Initialize command and config.
dependency.InitCommandAndConfig(rootCmd, true, cfg)
}
func initDfpath(cfg *config.ServerConfig) (dfpath.Dfpath, error) {
var options []dfpath.Option
if cfg.WorkHome != "" {
options = append(options, dfpath.WithWorkHome(cfg.WorkHome))
}
if cfg.LogDir != "" {
options = append(options, dfpath.WithLogDir(cfg.LogDir))
}
if cfg.CacheDir != "" {
options = append(options, dfpath.WithCacheDir(cfg.CacheDir))
}
if cfg.PluginDir != "" {
options = append(options, dfpath.WithPluginDir(cfg.PluginDir))
}
if cfg.DataDir != "" {
options = append(options, dfpath.WithDataDir(cfg.DataDir))
}
return dfpath.New(options...)
}
func runScheduler(ctx context.Context, d dfpath.Dfpath) error {
logger.Infof("version:\n%s", version.Version())
ff := dependency.InitMonitor(cfg.PProfPort, cfg.Telemetry)
defer ff()
svr, err := scheduler.New(ctx, cfg, d)
if err != nil {
return err
}
dependency.SetupQuitSignalHandler(func() { svr.Stop() })
return svr.Serve()
}