github.com/projecteru2/core@v0.0.0-20240321043226-06bcc1c23f58/store/etcdv3/processing.go (about)

     1  package etcdv3
     2  
     3  import (
     4  	"context"
     5  	"fmt"
     6  	"path/filepath"
     7  	"strconv"
     8  	"strings"
     9  
    10  	"github.com/projecteru2/core/log"
    11  	"github.com/projecteru2/core/types"
    12  
    13  	clientv3 "go.etcd.io/etcd/client/v3"
    14  )
    15  
    16  func (m *Mercury) getProcessingKey(processing *types.Processing) string {
    17  	return filepath.Join(workloadProcessingPrefix, processing.Appname, processing.Entryname, processing.Nodename, processing.Ident)
    18  }
    19  
    20  // CreateProcessing save processing status in etcd
    21  func (m *Mercury) CreateProcessing(ctx context.Context, processing *types.Processing, count int) error {
    22  	_, err := m.Create(ctx, m.getProcessingKey(processing), fmt.Sprintf("%d", count))
    23  	return err
    24  }
    25  
    26  // DeleteProcessing delete processing status in etcd
    27  func (m *Mercury) DeleteProcessing(ctx context.Context, processing *types.Processing) error {
    28  	_, err := m.Delete(ctx, m.getProcessingKey(processing))
    29  	return err
    30  }
    31  
    32  func (m *Mercury) doLoadProcessing(ctx context.Context, appname, entryname string) (map[string]int, error) {
    33  	nodesCount := map[string]int{}
    34  	// 显式地加 / 保证 prefix 一致性
    35  	processingKey := filepath.Join(workloadProcessingPrefix, appname, entryname) + "/"
    36  	resp, err := m.Get(ctx, processingKey, clientv3.WithPrefix())
    37  	if err != nil {
    38  		return nil, err
    39  	}
    40  	if resp.Count == 0 {
    41  		return nodesCount, nil
    42  	}
    43  	logger := log.WithFunc("store.etcdv3.doLoadProcessing")
    44  
    45  	for _, ev := range resp.Kvs {
    46  		key := string(ev.Key)
    47  		parts := strings.Split(key, "/")
    48  		nodename := parts[len(parts)-2]
    49  		count, err := strconv.Atoi(string(ev.Value))
    50  		if err != nil {
    51  			logger.Error(ctx, err, "Load processing status failed")
    52  			continue
    53  		}
    54  		if _, ok := nodesCount[nodename]; !ok {
    55  			nodesCount[nodename] = count
    56  			continue
    57  		}
    58  		nodesCount[nodename] += count
    59  	}
    60  	logger.Debugf(ctx, "Processing result: %+v", nodesCount)
    61  	return nodesCount, nil
    62  }