以下是优化后的代码:
func (srv *Service) SNMPCollectDatas(ctx context.Context, rows []*MibCollectCfg) ([]*MibCollectInfo, error) {
var infos []*MibCollectInfo
t1 := time.Now()
fmt.Println("开始采集")
// create SNMP client pool
pool, err := CreateSNMPClientPool(rows)
if err != nil {
return nil, err
}
defer pool.Close()
// iterate over rows and collect data in parallel
for _, row := range rows {
ip := row.IP
community := row.Community
snOid := row.SnOid
info := &MibCollectInfo{
IP: ip,
Community: community,
}
infos = append(infos, info)
eg := errgroup.Group{}
eg.SetLimit(100)
// collect SN data
eg.Go(func() error {
value, err := CollectSNData(pool, ip, community, snOid)
if err != nil {
log.Println(err.Error())
return nil
}
info.SN = value
return nil
})
// collect node data
nodes := getDefaultNodes()
for _, node := range nodes {
oid := node.Oid
name := node.Name
eg.Go(func() error {
nodeInfos, err := CollectNodeData(pool, ip, community, oid, name)
if err != nil {
nodeInfos = append(nodeInfos, &MibCollectData{
Oid: oid,
Name: name,
ErrMsg: err.Error(),
})
}
info.Rows = append(info.Rows, nodeInfos...)
return nil
})
}
if err := eg.Wait(); err != nil {
return nil, err
}
}
fmt.Println("采集完成", time.Since(t1).String())
return infos, nil
}
// CreateSNMPClientPool creates a pool of SNMP clients for parallelism. func CreateSNMPClientPool(rows []*MibCollectCfg) (*gosnmp.Pool, error) {
var clients []*gosnmp.GoSNMP
for _, row := range rows {
client, err := getSNMPClient(row.IP, row.Community)
if err != nil {
return nil, err
}
defer client.Close()
clients = append(clients, client)
}
return gosnmp.NewPool(clients...)
}
// CollectSNData collects the SN data from an SNMP client. func CollectSNData(pool *gosnmp.Pool, ip, community, snOid string) (string, error) {
client, err := pool.Get(ip)
if err != nil {
return "", err
}
defer pool.Put(client)
valueUnit, err := client.Get(snOid)
if err != nil {
return "", err
}
return convertor.ToString(valueUnit.Value), nil
}
// CollectNodeData collects the node data from an SNMP client. func CollectNodeData(pool *gosnmp.Pool, ip, community, oid, name string) ([]*MibCollectData, error) {
client, err := pool.Get(ip)
if err != nil {
return nil, err
}
defer pool.Put(client)
var nodeInfos []*MibCollectData
err = client.BulkWalk(oid, func(valueUnit gosnmp.SnmpPDU) error {
tmpNodeInfo := &MibCollectData{
Oid: valueUnit.Name,
Name: name,
}
switch valueUnit.Type {
case gosnmp.IPAddress:
tmpNodeInfo.Value = net.ParseIP(string(valueUnit.Value.([]byte))).String()
case gosnmp.OctetString:
mac, err := net.ParseMAC(net.HardwareAddr(valueUnit.Value.([]byte)).String())
if err == nil {
tmpNodeInfo.Value = mac.String()
} else {
tmpNodeInfo.Value = string(valueUnit.Value.([]byte))
}
default:
tmpNodeInfo.Value = convertor.ToString(valueUnit.Value)
}
nodeInfos = append(nodeInfos, tmpNodeInfo)
return nil
})
return nodeInfos, err
}
优化的主要思路是:
- 使用 SNMP client 池来提高并发性能;
- 将 SN 和节点数据采集拆分为两个函数,以便更好地复用代码和测试;
- 使用 errgroup 控制并发,避免同时向每个设备发送太多请求造成网络阻塞;
- 增加错误处理,确保程序在出现错误时能够正确处理,并返回错误信息。