crawlab/core/node/service/master_service.go

package service

import (
	"github.com/apex/log"
	config2 "github.com/crawlab-team/crawlab/core/config"
	"github.com/crawlab-team/crawlab/core/constants"
	"github.com/crawlab-team/crawlab/core/container"
	"github.com/crawlab-team/crawlab/core/errors"
	"github.com/crawlab-team/crawlab/core/grpc/server"
	"github.com/crawlab-team/crawlab/core/interfaces"
	"github.com/crawlab-team/crawlab/core/models/common"
	"github.com/crawlab-team/crawlab/core/models/delegate"
	"github.com/crawlab-team/crawlab/core/models/models"
	"github.com/crawlab-team/crawlab/core/models/service"
	"github.com/crawlab-team/crawlab/core/node/config"
	"github.com/crawlab-team/crawlab/core/notification"
	"github.com/crawlab-team/crawlab/core/system"
	"github.com/crawlab-team/crawlab/core/utils"
	"github.com/crawlab-team/crawlab/db/mongo"
	grpc "github.com/crawlab-team/crawlab/grpc"
	"github.com/crawlab-team/go-trace"
	"github.com/spf13/viper"
	"go.mongodb.org/mongo-driver/bson"
	mongo2 "go.mongodb.org/mongo-driver/mongo"
	"time"
)

type MasterService struct {
	// dependencies
	modelSvc        service.ModelService
	cfgSvc          interfaces.NodeConfigService
	server          interfaces.GrpcServer
	schedulerSvc    interfaces.TaskSchedulerService
	handlerSvc      interfaces.TaskHandlerService
	scheduleSvc     interfaces.ScheduleService
	notificationSvc *notification.Service
	spiderAdminSvc  interfaces.SpiderAdminService
	systemSvc       *system.Service

	// settings
	cfgPath         string
	address         interfaces.Address
	monitorInterval time.Duration
	stopOnError     bool
}

func (svc *MasterService) Init() (err error) {
	// do nothing
	return nil
}

func (svc *MasterService) Start() {
	// create indexes
	common.CreateIndexes()

	// start grpc server
	if err := svc.server.Start(); err != nil {
		panic(err)
	}

	// register to db
	if err := svc.Register(); err != nil {
		panic(err)
	}

	// start monitoring worker nodes
	go svc.Monitor()

	// start task handler
	go svc.handlerSvc.Start()

	// start task scheduler
	go svc.schedulerSvc.Start()

	// start schedule service
	go svc.scheduleSvc.Start()

	// start notification service
	go svc.notificationSvc.Start()

	// start spider admin service
	go svc.spiderAdminSvc.Start()

	// wait for quit signal
	svc.Wait()

	// stop
	svc.Stop()
}

func (svc *MasterService) Wait() {
	utils.DefaultWait()
}

func (svc *MasterService) Stop() {
	_ = svc.server.Stop()
	log.Infof("master[%s] service has stopped", svc.GetConfigService().GetNodeKey())
}

func (svc *MasterService) Monitor() {
	log.Infof("master[%s] monitoring started", svc.GetConfigService().GetNodeKey())
	for {
		if err := svc.monitor(); err != nil {
			trace.PrintError(err)
			if svc.stopOnError {
				log.Errorf("master[%s] monitor error, now stopping...", svc.GetConfigService().GetNodeKey())
				svc.Stop()
				return
			}
		}

		time.Sleep(svc.monitorInterval)
	}
}

func (svc *MasterService) GetConfigService() (cfgSvc interfaces.NodeConfigService) {
	return svc.cfgSvc
}

func (svc *MasterService) GetConfigPath() (path string) {
	return svc.cfgPath
}

func (svc *MasterService) SetConfigPath(path string) {
	svc.cfgPath = path
}

func (svc *MasterService) GetAddress() (address interfaces.Address) {
	return svc.address
}

func (svc *MasterService) SetAddress(address interfaces.Address) {
	svc.address = address
}

func (svc *MasterService) SetMonitorInterval(duration time.Duration) {
	svc.monitorInterval = duration
}

func (svc *MasterService) Register() (err error) {
	nodeKey := svc.GetConfigService().GetNodeKey()
	nodeName := svc.GetConfigService().GetNodeName()
	node, err := svc.modelSvc.GetNodeByKey(nodeKey, nil)
	if err != nil && err.Error() == mongo2.ErrNoDocuments.Error() {
		// not exists
		log.Infof("master[%s] does not exist in db", nodeKey)
		node := &models.Node{
			Key:        nodeKey,
			Name:       nodeName,
			MaxRunners: config.DefaultConfigOptions.MaxRunners,
			IsMaster:   true,
			Status:     constants.NodeStatusOnline,
			Enabled:    true,
			Active:     true,
			ActiveTs:   time.Now(),
		}
		if viper.GetInt("task.handler.maxRunners") > 0 {
			node.MaxRunners = viper.GetInt("task.handler.maxRunners")
		}
		nodeD := delegate.NewModelNodeDelegate(node)
		if err := nodeD.Add(); err != nil {
			return err
		}
		log.Infof("added master[%s] in db. id: %s", nodeKey, nodeD.GetModel().GetId().Hex())
		return nil
	} else if err == nil {
		// exists
		log.Infof("master[%s] exists in db", nodeKey)
		nodeD := delegate.NewModelNodeDelegate(node)
		if err := nodeD.UpdateStatusOnline(); err != nil {
			return err
		}
		log.Infof("updated master[%s] in db. id: %s", nodeKey, nodeD.GetModel().GetId().Hex())
		return nil
	} else {
		// error
		return err
	}
}

func (svc *MasterService) StopOnError() {
	svc.stopOnError = true
}

func (svc *MasterService) GetServer() (svr interfaces.GrpcServer) {
	return svc.server
}

func (svc *MasterService) monitor() (err error) {
	// update master node status in db
	if err := svc.updateMasterNodeStatus(); err != nil {
		if err.Error() == mongo2.ErrNoDocuments.Error() {
			return nil
		}
		return err
	}

	// all worker nodes
	nodes, err := svc.getAllWorkerNodes()
	if err != nil {
		return err
	}

	// error flag
	isErr := false

	// iterate all nodes
	for _, n := range nodes {
		// subscribe
		if err := svc.subscribeNode(&n); err != nil {
			isErr = true
			continue
		}

		// ping client
		if err := svc.pingNodeClient(&n); err != nil {
			isErr = true
			continue
		}

		// update node available runners
		if err := svc.updateNodeAvailableRunners(&n); err != nil {
			isErr = true
			continue
		}
	}

	if isErr {
		return trace.TraceError(errors.ErrorNodeMonitorError)
	}

	return nil
}

func (svc *MasterService) getAllWorkerNodes() (nodes []models.Node, err error) {
	query := bson.M{
		"key":    bson.M{"$ne": svc.cfgSvc.GetNodeKey()}, // not self
		"active": true,                                   // active
	}
	nodes, err = svc.modelSvc.GetNodeList(query, nil)
	if err != nil {
		if err == mongo2.ErrNoDocuments {
			return nil, nil
		}
		return nil, trace.TraceError(err)
	}
	return nodes, nil
}

func (svc *MasterService) updateMasterNodeStatus() (err error) {
	nodeKey := svc.GetConfigService().GetNodeKey()
	node, err := svc.modelSvc.GetNodeByKey(nodeKey, nil)
	if err != nil {
		return err
	}
	nodeD := delegate.NewModelNodeDelegate(node)
	return nodeD.UpdateStatusOnline()
}

func (svc *MasterService) setWorkerNodeOffline(n interfaces.Node) (err error) {
	return delegate.NewModelNodeDelegate(n).UpdateStatusOffline()
}

func (svc *MasterService) subscribeNode(n interfaces.Node) (err error) {
	_, err = svc.server.GetSubscribe("node:" + n.GetKey())
	if err != nil {
		log.Errorf("cannot subscribe worker node[%s]: %v", n.GetKey(), err)
		if err := svc.setWorkerNodeOffline(n); err != nil {
			return trace.TraceError(err)
		}
		return trace.TraceError(err)
	}
	return nil
}

func (svc *MasterService) pingNodeClient(n interfaces.Node) (err error) {
	if err := svc.server.SendStreamMessage("node:"+n.GetKey(), grpc.StreamMessageCode_PING); err != nil {
		log.Errorf("cannot ping worker node client[%s]: %v", n.GetKey(), err)
		if err := svc.setWorkerNodeOffline(n); err != nil {
			return trace.TraceError(err)
		}
		return trace.TraceError(err)
	}
	return nil
}

func (svc *MasterService) updateNodeAvailableRunners(n interfaces.Node) (err error) {
	query := bson.M{
		"node_id": n.GetId(),
		"status":  constants.TaskStatusRunning,
	}
	runningTasksCount, err := mongo.GetMongoCol(interfaces.ModelColNameTask).Count(query)
	if err != nil {
		return trace.TraceError(err)
	}
	n.SetAvailableRunners(n.GetMaxRunners() - runningTasksCount)
	return delegate.NewModelDelegate(n).Save()
}

func NewMasterService(opts ...Option) (res interfaces.NodeMasterService, err error) {
	// master service
	svc := &MasterService{
		cfgPath:         config2.GetConfigPath(),
		monitorInterval: 15 * time.Second,
		stopOnError:     false,
	}

	// apply options
	for _, opt := range opts {
		opt(svc)
	}

	// server options
	var serverOpts []server.Option
	if svc.address != nil {
		serverOpts = append(serverOpts, server.WithAddress(svc.address))
	}

	// dependency injection
	if err := container.GetContainer().Invoke(func(
		cfgSvc interfaces.NodeConfigService,
		modelSvc service.ModelService,
		server interfaces.GrpcServer,
		schedulerSvc interfaces.TaskSchedulerService,
		handlerSvc interfaces.TaskHandlerService,
		scheduleSvc interfaces.ScheduleService,
		spiderAdminSvc interfaces.SpiderAdminService,
	) {
		svc.cfgSvc = cfgSvc
		svc.modelSvc = modelSvc
		svc.server = server
		svc.schedulerSvc = schedulerSvc
		svc.handlerSvc = handlerSvc
		svc.scheduleSvc = scheduleSvc
		svc.spiderAdminSvc = spiderAdminSvc
	}); err != nil {
		return nil, err
	}

	// notification service
	svc.notificationSvc = notification.GetService()

	// system service
	svc.systemSvc = system.GetService()

	// init
	if err := svc.Init(); err != nil {
		return nil, err
	}

	return svc, nil
}

func ProvideMasterService(path string, opts ...Option) func() (interfaces.NodeMasterService, error) {
	// path
	if path == "" || path == config2.GetConfigPath() {
		if viper.GetString("config.path") != "" {
			path = viper.GetString("config.path")
		} else {
			path = config2.GetConfigPath()
		}
	}
	opts = append(opts, WithConfigPath(path))

	return func() (interfaces.NodeMasterService, error) {
		return NewMasterService(opts...)
	}
}