dragonfly/scheduler/core/evaluator/evaluator.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 evaluator

import (
	"sort"
	"strings"
	"sync"

	"d7y.io/dragonfly/v2/internal/idgen"
	"d7y.io/dragonfly/v2/scheduler/config"
	"d7y.io/dragonfly/v2/scheduler/supervisor"
)

type Evaluator interface {

	// Evaluate todo Normalization
	Evaluate(parent *supervisor.Peer, child *supervisor.Peer) float64

	// NeedAdjustParent determine whether the peer needs a new parent node
	NeedAdjustParent(peer *supervisor.Peer) bool

	// IsBadNode determine if peer is a failed node
	IsBadNode(peer *supervisor.Peer) bool
}

type Factory struct {
	lock                         sync.RWMutex
	evaluators                   map[string]Evaluator
	getEvaluatorFuncs            map[int]getEvaluatorFunc
	getEvaluatorFuncPriorityList []getEvaluatorFunc
	cache                        map[string]Evaluator
	cacheClearFunc               sync.Once
	abtest                       bool
	aEvaluator                   string
	bEvaluator                   string
}

var _ Evaluator = (*Factory)(nil)

func (ef *Factory) Evaluate(dst *supervisor.Peer, src *supervisor.Peer) float64 {
	return ef.get(dst.Task.TaskID).Evaluate(dst, src)
}

func (ef *Factory) NeedAdjustParent(peer *supervisor.Peer) bool {
	return ef.get(peer.Task.TaskID).NeedAdjustParent(peer)
}

func (ef *Factory) IsBadNode(peer *supervisor.Peer) bool {
	return ef.get(peer.Task.TaskID).IsBadNode(peer)
}

func NewEvaluatorFactory(cfg *config.SchedulerConfig) *Factory {
	factory := &Factory{
		evaluators:        make(map[string]Evaluator),
		getEvaluatorFuncs: map[int]getEvaluatorFunc{},
		cache:             map[string]Evaluator{},
		abtest:            cfg.ABTest,
		aEvaluator:        cfg.AEvaluator,
		bEvaluator:        cfg.BEvaluator,
	}
	return factory
}

var (
	m = make(map[string]Evaluator)
)

func Register(name string, evaluator Evaluator) {
	m[strings.ToLower(name)] = evaluator
}

func Get(name string) Evaluator {
	if eval, ok := m[strings.ToLower(name)]; ok {
		return eval
	}
	return nil
}

type getEvaluatorFunc func(taskID string) (string, bool)

func (ef *Factory) get(taskID string) Evaluator {
	ef.lock.RLock()
	evaluator, ok := ef.cache[taskID]
	ef.lock.RUnlock()
	if ok {
		return evaluator
	}

	if ef.abtest {
		name := ""
		if strings.HasSuffix(taskID, idgen.TwinsBSuffix) {
			if ef.bEvaluator != "" {
				name = ef.bEvaluator
			}
		} else {
			if ef.aEvaluator != "" {
				name = ef.aEvaluator
			}
		}
		if name != "" {
			ef.lock.RLock()
			evaluator, ok = ef.evaluators[name]
			ef.lock.RUnlock()
			if ok {
				ef.lock.Lock()
				ef.cache[taskID] = evaluator
				ef.lock.Unlock()
				return evaluator
			}
		}
	}

	for _, fun := range ef.getEvaluatorFuncPriorityList {
		name, ok := fun(taskID)
		if !ok {
			continue
		}
		ef.lock.RLock()
		evaluator, ok = ef.evaluators[name]
		ef.lock.RUnlock()
		if !ok {
			continue
		}

		ef.lock.Lock()
		ef.cache[taskID] = evaluator
		ef.lock.Unlock()
		return evaluator
	}

	return nil
}

func (ef *Factory) clearCache() {
	ef.lock.Lock()
	ef.cache = make(map[string]Evaluator)
	ef.lock.Unlock()
}

func (ef *Factory) add(name string, evaluator Evaluator) {
	ef.lock.Lock()
	ef.evaluators[name] = evaluator
	ef.lock.Unlock()
}

func (ef *Factory) addGetEvaluatorFunc(priority int, fun getEvaluatorFunc) {
	ef.lock.Lock()
	defer ef.lock.Unlock()
	_, ok := ef.getEvaluatorFuncs[priority]
	if ok {
		return
	}
	ef.getEvaluatorFuncs[priority] = fun
	var priorities []int
	for p := range ef.getEvaluatorFuncs {
		priorities = append(priorities, p)
	}
	sort.Ints(priorities)
	ef.getEvaluatorFuncPriorityList = ef.getEvaluatorFuncPriorityList[:0]
	for i := len(priorities) - 1; i >= 0; i-- {
		ef.getEvaluatorFuncPriorityList = append(ef.getEvaluatorFuncPriorityList, ef.getEvaluatorFuncs[priorities[i]])
	}

}

func (ef *Factory) deleteGetEvaluatorFunc(priority int, fun getEvaluatorFunc) {
	ef.lock.Lock()

	delete(ef.getEvaluatorFuncs, priority)

	var priorities []int
	for p := range ef.getEvaluatorFuncs {
		priorities = append(priorities, p)
	}
	sort.Ints(priorities)
	ef.getEvaluatorFuncPriorityList = ef.getEvaluatorFuncPriorityList[:0]
	for i := len(priorities) - 1; i >= 0; i-- {
		ef.getEvaluatorFuncPriorityList = append(ef.getEvaluatorFuncPriorityList, ef.getEvaluatorFuncs[priorities[i]])
	}

	ef.lock.Unlock()
}

func (ef *Factory) Register(name string, evaluator Evaluator) {
	//ef.cacheClearFunc.Do(func() {
	//	tick := time.NewTicker(time.Hour)
	//	for {
	//		select {
	//		case <-tick.C:
	//			ef.clearCache()
	//		}
	//	}
	//})
	ef.add(name, evaluator)
	ef.clearCache()
}

func (ef *Factory) RegisterGetEvaluatorFunc(priority int, fun getEvaluatorFunc) {
	ef.addGetEvaluatorFunc(priority, fun)
	ef.clearCache()
}