github.com/TeaOSLab/EdgeNode@v1.3.8/internal/metrics/manager.go (about)

     1  // Copyright 2021 Liuxiangchao iwind.liu@gmail.com. All rights reserved.
     2  
     3  package metrics
     4  
     5  import (
     6  	"github.com/TeaOSLab/EdgeCommon/pkg/serverconfigs"
     7  	teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
     8  	"github.com/TeaOSLab/EdgeNode/internal/events"
     9  	"github.com/TeaOSLab/EdgeNode/internal/remotelogs"
    10  	"strconv"
    11  	"sync"
    12  )
    13  
    14  var SharedManager = NewManager()
    15  
    16  func init() {
    17  	if !teaconst.IsMain {
    18  		return
    19  	}
    20  
    21  	events.OnClose(func() {
    22  		SharedManager.Quit()
    23  	})
    24  }
    25  
    26  type Manager struct {
    27  	isQuiting bool
    28  
    29  	taskMap         map[int64]Task    // itemId => Task
    30  	categoryTaskMap map[string][]Task // category => []Task
    31  	locker          sync.RWMutex
    32  
    33  	hasHTTPMetrics bool
    34  	hasTCPMetrics  bool
    35  	hasUDPMetrics  bool
    36  }
    37  
    38  func NewManager() *Manager {
    39  	return &Manager{
    40  		taskMap:         map[int64]Task{},
    41  		categoryTaskMap: map[string][]Task{},
    42  	}
    43  }
    44  
    45  func (this *Manager) Update(items []*serverconfigs.MetricItemConfig) {
    46  	if this.isQuiting {
    47  		return
    48  	}
    49  
    50  	this.locker.Lock()
    51  	defer this.locker.Unlock()
    52  
    53  	var newMap = map[int64]*serverconfigs.MetricItemConfig{}
    54  	for _, item := range items {
    55  		newMap[item.Id] = item
    56  	}
    57  
    58  	// 停用以前的 或 修改现在的
    59  	for itemId, task := range this.taskMap {
    60  		newItem, ok := newMap[itemId]
    61  		if !ok || !newItem.IsOn { // 停用以前的
    62  			remotelogs.Println("METRIC_MANAGER", "stop task '"+strconv.FormatInt(itemId, 10)+"'")
    63  			err := task.Stop()
    64  			if err != nil {
    65  				remotelogs.Error("METRIC_MANAGER", "stop task '"+strconv.FormatInt(itemId, 10)+"' failed: "+err.Error())
    66  			}
    67  
    68  			// deleted
    69  			if newItem != nil && !newItem.IsOn {
    70  				deleteErr := task.Delete()
    71  				if deleteErr != nil {
    72  					remotelogs.Error("METRIC_MANAGER", "delete task '"+strconv.FormatInt(itemId, 10)+"' failed: "+err.Error())
    73  				}
    74  			}
    75  
    76  			delete(this.taskMap, itemId)
    77  		} else { // 更新已存在的
    78  			if newItem.Version != task.Item().Version {
    79  				remotelogs.Println("METRIC_MANAGER", "update task '"+strconv.FormatInt(itemId, 10)+"'")
    80  				task.SetItem(newItem)
    81  			}
    82  		}
    83  	}
    84  
    85  	// 启动新的
    86  	for _, newItem := range items {
    87  		if !newItem.IsOn {
    88  			continue
    89  		}
    90  		_, ok := this.taskMap[newItem.Id]
    91  		if !ok {
    92  			remotelogs.Println("METRIC_MANAGER", "start task '"+strconv.FormatInt(newItem.Id, 10)+"'")
    93  			var task Task
    94  
    95  			if CheckSQLiteDB(newItem.Id) || !teaconst.EnableKVCacheStore {
    96  				task = NewSQLiteTask(newItem)
    97  			} else {
    98  				task = NewKVTask(newItem)
    99  			}
   100  
   101  			err := task.Init()
   102  			if err != nil {
   103  				remotelogs.Error("METRIC_MANAGER", "initialized task failed: "+err.Error())
   104  				continue
   105  			}
   106  			err = task.Start()
   107  			if err != nil {
   108  				remotelogs.Error("METRIC_MANAGER", "start task failed: "+err.Error())
   109  				continue
   110  			}
   111  			this.taskMap[newItem.Id] = task
   112  		}
   113  	}
   114  
   115  	// 按分类存放
   116  	this.hasHTTPMetrics = false
   117  	this.hasTCPMetrics = false
   118  	this.hasUDPMetrics = false
   119  	this.categoryTaskMap = map[string][]Task{}
   120  	for _, task := range this.taskMap {
   121  		var tasks = this.categoryTaskMap[task.Item().Category]
   122  		tasks = append(tasks, task)
   123  		this.categoryTaskMap[task.Item().Category] = tasks
   124  
   125  		switch task.Item().Category {
   126  		case serverconfigs.MetricItemCategoryHTTP:
   127  			this.hasHTTPMetrics = true
   128  		case serverconfigs.MetricItemCategoryTCP:
   129  			this.hasTCPMetrics = true
   130  		case serverconfigs.MetricItemCategoryUDP:
   131  			this.hasUDPMetrics = true
   132  		}
   133  	}
   134  }
   135  
   136  // Add 添加数据
   137  func (this *Manager) Add(obj MetricInterface) {
   138  	if this.isQuiting {
   139  		return
   140  	}
   141  
   142  	this.locker.RLock()
   143  	var tasks = this.categoryTaskMap[obj.MetricCategory()]
   144  	this.locker.RUnlock()
   145  
   146  	for _, task := range tasks {
   147  		task.Add(obj)
   148  	}
   149  }
   150  
   151  func (this *Manager) HasHTTPMetrics() bool {
   152  	return this.hasHTTPMetrics
   153  }
   154  
   155  func (this *Manager) HasTCPMetrics() bool {
   156  	return this.hasTCPMetrics
   157  }
   158  
   159  func (this *Manager) HasUDPMetrics() bool {
   160  	return this.hasUDPMetrics
   161  }
   162  
   163  func (this *Manager) TaskMap() map[int64]Task {
   164  	return this.taskMap
   165  }
   166  
   167  // Quit 退出管理器
   168  func (this *Manager) Quit() {
   169  	this.isQuiting = true
   170  
   171  	remotelogs.Println("METRIC_MANAGER", "quit")
   172  
   173  	this.locker.Lock()
   174  	for _, task := range this.taskMap {
   175  		_ = task.Stop()
   176  	}
   177  	this.taskMap = map[int64]Task{}
   178  	this.locker.Unlock()
   179  }