github.com/TeaOSLab/EdgeNode@v1.3.8/internal/nodes/task_sync_api_nodes.go (about)

     1  package nodes
     2  
     3  import (
     4  	"github.com/TeaOSLab/EdgeCommon/pkg/rpc/pb"
     5  	"github.com/TeaOSLab/EdgeNode/internal/configs"
     6  	teaconst "github.com/TeaOSLab/EdgeNode/internal/const"
     7  	"github.com/TeaOSLab/EdgeNode/internal/events"
     8  	"github.com/TeaOSLab/EdgeNode/internal/goman"
     9  	"github.com/TeaOSLab/EdgeNode/internal/rpc"
    10  	"github.com/TeaOSLab/EdgeNode/internal/trackers"
    11  	"github.com/TeaOSLab/EdgeNode/internal/utils"
    12  	"github.com/iwind/TeaGo/Tea"
    13  	"github.com/iwind/TeaGo/logs"
    14  	"time"
    15  )
    16  
    17  var sharedSyncAPINodesTask = NewSyncAPINodesTask()
    18  
    19  func init() {
    20  	if !teaconst.IsMain {
    21  		return
    22  	}
    23  
    24  	events.On(events.EventStart, func() {
    25  		goman.New(func() {
    26  			sharedSyncAPINodesTask.Start()
    27  		})
    28  	})
    29  	events.OnClose(func() {
    30  		sharedSyncAPINodesTask.Stop()
    31  	})
    32  }
    33  
    34  // SyncAPINodesTask API节点同步任务
    35  type SyncAPINodesTask struct {
    36  	ticker *time.Ticker
    37  }
    38  
    39  func NewSyncAPINodesTask() *SyncAPINodesTask {
    40  	return &SyncAPINodesTask{}
    41  }
    42  
    43  func (this *SyncAPINodesTask) Start() {
    44  	this.ticker = time.NewTicker(5 * time.Minute)
    45  	if Tea.IsTesting() {
    46  		// 快速测试
    47  		this.ticker = time.NewTicker(1 * time.Minute)
    48  	}
    49  	for range this.ticker.C {
    50  		err := this.Loop()
    51  		if err != nil {
    52  			logs.Println("[TASK][SYNC_API_NODES_TASK]" + err.Error())
    53  		}
    54  	}
    55  }
    56  
    57  func (this *SyncAPINodesTask) Stop() {
    58  	if this.ticker != nil {
    59  		this.ticker.Stop()
    60  	}
    61  }
    62  
    63  func (this *SyncAPINodesTask) Loop() error {
    64  	// 如果有节点定制的API节点地址
    65  	var hasCustomizedAPINodeAddrs = sharedNodeConfig != nil && len(sharedNodeConfig.APINodeAddrs) > 0
    66  
    67  	config, err := configs.LoadAPIConfig()
    68  	if err != nil {
    69  		return err
    70  	}
    71  
    72  	// 是否禁止自动升级
    73  	if config.RPCDisableUpdate {
    74  		return nil
    75  	}
    76  
    77  	var tr = trackers.Begin("SYNC_API_NODES")
    78  	defer tr.End()
    79  
    80  	// 获取所有可用的节点
    81  	rpcClient, err := rpc.SharedRPC()
    82  	if err != nil {
    83  		return err
    84  	}
    85  	resp, err := rpcClient.APINodeRPC.FindAllEnabledAPINodes(rpcClient.Context(), &pb.FindAllEnabledAPINodesRequest{})
    86  	if err != nil {
    87  		return err
    88  	}
    89  
    90  	var newEndpoints = []string{}
    91  	for _, node := range resp.ApiNodes {
    92  		if !node.IsOn {
    93  			continue
    94  		}
    95  		newEndpoints = append(newEndpoints, node.AccessAddrs...)
    96  	}
    97  
    98  	// 和现有的对比
    99  	if utils.EqualStrings(newEndpoints, config.RPCEndpoints) {
   100  		return nil
   101  	}
   102  
   103  	// 测试是否有API节点可用
   104  	var hasOk = rpcClient.TestEndpoints(newEndpoints)
   105  	if !hasOk {
   106  		return nil
   107  	}
   108  
   109  	// 修改RPC对象配置
   110  	config.RPCEndpoints = newEndpoints
   111  
   112  	// 更新当前RPC
   113  	if !hasCustomizedAPINodeAddrs {
   114  		err = rpcClient.UpdateConfig(config)
   115  		if err != nil {
   116  			return err
   117  		}
   118  	}
   119  
   120  	// 保存到文件
   121  	err = config.WriteFile(Tea.ConfigFile(configs.ConfigFileName))
   122  	if err != nil {
   123  		return err
   124  	}
   125  
   126  	return nil
   127  }