github.com/projecteru2/core@v0.0.0-20240321043226-06bcc1c23f58/store/redis/workload.go

package redis

import (
	"context"
	"encoding/json"
	"fmt"
	"path/filepath"

	"github.com/projecteru2/core/log"
	"github.com/projecteru2/core/types"
	"github.com/projecteru2/core/utils"
)

// AddWorkload add a workload
// mainly record its relationship on pod and node
// actually if we already know its node, we will know its pod
// but we still store it
// storage path in etcd is `/workload/:workloadid`
func (r *Rediaron) AddWorkload(ctx context.Context, workload *types.Workload, processing *types.Processing) error {
	return r.doOpsWorkload(ctx, workload, processing, true)
}

// UpdateWorkload update a workload
func (r *Rediaron) UpdateWorkload(ctx context.Context, workload *types.Workload) error {
	return r.doOpsWorkload(ctx, workload, nil, false)
}

// RemoveWorkload remove a workload
// workload ID must be in full length
func (r *Rediaron) RemoveWorkload(ctx context.Context, workload *types.Workload) error {
	return r.cleanWorkloadData(ctx, workload)
}

// GetWorkload get a workload
// workload if must be in full length, or we can't find it in etcd
// storage path in etcd is `/workload/:workloadid`
func (r *Rediaron) GetWorkload(ctx context.Context, ID string) (*types.Workload, error) {
	workloads, err := r.GetWorkloads(ctx, []string{ID})
	if err != nil {
		return nil, err
	}
	return workloads[0], nil
}

// GetWorkloads get many workloads
func (r *Rediaron) GetWorkloads(ctx context.Context, IDs []string) (workloads []*types.Workload, err error) {
	keys := []string{}
	for _, ID := range IDs {
		keys = append(keys, fmt.Sprintf(workloadInfoKey, ID))
	}

	return r.doGetWorkloads(ctx, keys)
}

// GetWorkloadStatus get workload status
func (r *Rediaron) GetWorkloadStatus(ctx context.Context, ID string) (*types.StatusMeta, error) {
	workload, err := r.GetWorkload(ctx, ID)
	if err != nil {
		return nil, err
	}
	return workload.StatusMeta, nil
}

// SetWorkloadStatus set workload status
func (r *Rediaron) SetWorkloadStatus(ctx context.Context, status *types.StatusMeta, ttl int64) error {
	if status.Appname == "" || status.Entrypoint == "" || status.Nodename == "" {
		return types.ErrInvaildWorkloadStatus
	}

	data, err := json.Marshal(status)
	if err != nil {
		return err
	}
	statusVal := string(data)
	statusKey := filepath.Join(workloadStatusPrefix, status.Appname, status.Entrypoint, status.Nodename, status.ID)
	workloadKey := fmt.Sprintf(workloadInfoKey, status.ID)
	return r.BindStatus(ctx, workloadKey, statusKey, statusVal, ttl)
}

// ListWorkloads list workloads
func (r *Rediaron) ListWorkloads(ctx context.Context, appname, entrypoint, nodename string, limit int64, labels map[string]string) ([]*types.Workload, error) {
	if appname == "" {
		entrypoint = ""
	}
	if entrypoint == "" {
		nodename = ""
	}
	// 这里显式加个 / 来保证 prefix 是唯一的
	key := filepath.Join(workloadDeployPrefix, appname, entrypoint, nodename) + "/*"
	data, err := r.getByKeyPattern(ctx, key, limit)
	if err != nil {
		return nil, err
	}

	workloads := []*types.Workload{}
	for _, v := range data {
		workload := &types.Workload{}
		if err := json.Unmarshal([]byte(v), workload); err != nil {
			return nil, err
		}
		if utils.LabelsFilter(workload.Labels, labels) {
			workloads = append(workloads, workload)
		}
	}

	return r.bindWorkloadsAdditions(ctx, workloads)
}

// ListNodeWorkloads list workloads belong to one node
func (r *Rediaron) ListNodeWorkloads(ctx context.Context, nodename string, labels map[string]string) ([]*types.Workload, error) {
	key := fmt.Sprintf(nodeWorkloadsKey, nodename, "*")
	data, err := r.getByKeyPattern(ctx, key, 0)
	if err != nil {
		return nil, err
	}

	workloads := []*types.Workload{}
	for _, v := range data {
		workload := &types.Workload{}
		if err := json.Unmarshal([]byte(v), workload); err != nil {
			return nil, err
		}
		if utils.LabelsFilter(workload.Labels, labels) {
			workloads = append(workloads, workload)
		}
	}

	return r.bindWorkloadsAdditions(ctx, workloads)
}

// WorkloadStatusStream watch deployed status
func (r *Rediaron) WorkloadStatusStream(ctx context.Context, appname, entrypoint, nodename string, labels map[string]string) chan *types.WorkloadStatus {
	if appname == "" {
		entrypoint = ""
	}
	if entrypoint == "" {
		nodename = ""
	}
	// 显式加个 / 保证 prefix 唯一
	statusKey := filepath.Join(workloadStatusPrefix, appname, entrypoint, nodename) + "/*"
	ch := make(chan *types.WorkloadStatus)
	logger := log.WithFunc("store.redis.WorkloadStatusStream")
	_ = r.pool.Invoke(func() {
		defer func() {
			logger.Info(ctx, "close WorkloadStatus channel")
			close(ch)
		}()

		logger.Infof(ctx, "watch on %s", statusKey)
		for message := range r.KNotify(ctx, statusKey) {
			_, _, _, ID := parseStatusKey(message.Key)
			msg := &types.WorkloadStatus{
				ID:     ID,
				Delete: message.Action == actionDel || message.Action == actionExpired,
			}
			workload, err := r.GetWorkload(ctx, ID)
			switch {
			case err != nil:
				msg.Error = err
			case utils.LabelsFilter(workload.Labels, labels):
				logger.Debugf(ctx, "workload %s status changed", workload.ID)
				msg.Workload = workload
			default:
				continue
			}
			ch <- msg
		}
	})
	return ch
}

func (r *Rediaron) cleanWorkloadData(ctx context.Context, workload *types.Workload) error {
	appname, entrypoint, _, err := utils.ParseWorkloadName(workload.Name)
	if err != nil {
		return err
	}

	keys := []string{
		filepath.Join(workloadStatusPrefix, appname, entrypoint, workload.Nodename, workload.ID), // workload deploy status
		filepath.Join(workloadDeployPrefix, appname, entrypoint, workload.Nodename, workload.ID), // workload deploy status
		fmt.Sprintf(workloadInfoKey, workload.ID),                                                // workload info
		fmt.Sprintf(nodeWorkloadsKey, workload.Nodename, workload.ID),                            // node workloads
	}
	return r.BatchDelete(ctx, keys)
}

func (r *Rediaron) doGetWorkloads(ctx context.Context, keys []string) ([]*types.Workload, error) {
	data, err := r.GetMulti(ctx, keys)
	if err != nil {
		return nil, err
	}

	workloads := []*types.Workload{}
	for k, v := range data {
		workload := &types.Workload{}
		if err = json.Unmarshal([]byte(v), workload); err != nil {
			log.WithFunc("store.redis.doGetWorkloads").Errorf(ctx, err, "failed to unmarshal %+v", k)
			return nil, err
		}
		workloads = append(workloads, workload)
	}

	return r.bindWorkloadsAdditions(ctx, workloads)
}

func (r *Rediaron) bindWorkloadsAdditions(ctx context.Context, workloads []*types.Workload) ([]*types.Workload, error) {
	nodes := map[string]*types.Node{}
	nodenames := []string{}
	nodenameCache := map[string]struct{}{}
	statusKeys := map[string]string{}
	logger := log.WithFunc("store.redis.bindWorkloadsAdditions")
	for _, workload := range workloads {
		appname, entrypoint, _, err := utils.ParseWorkloadName(workload.Name)
		if err != nil {
			return nil, err
		}
		statusKeys[workload.ID] = filepath.Join(workloadStatusPrefix, appname, entrypoint, workload.Nodename, workload.ID)
		if _, ok := nodenameCache[workload.Nodename]; !ok {
			nodenameCache[workload.Nodename] = struct{}{}
			nodenames = append(nodenames, workload.Nodename)
		}
	}
	ns, err := r.GetNodes(ctx, nodenames)
	if err != nil {
		return nil, err
	}
	for _, node := range ns {
		nodes[node.Name] = node
	}

	for index, workload := range workloads {
		if _, ok := nodes[workload.Nodename]; !ok {
			return nil, types.ErrInvaildWorkloadMeta
		}
		workloads[index].Engine = nodes[workload.Nodename].Engine
		if _, ok := statusKeys[workload.ID]; !ok {
			continue
		}
		v, err := r.GetOne(ctx, statusKeys[workload.ID])
		if err != nil {
			continue
		}
		status := &types.StatusMeta{}
		if err := json.Unmarshal([]byte(v), &status); err != nil {
			logger.Warnf(ctx, "unmarshal %s status data, raw %s", workload.ID, v)
			logger.Error(ctx, err, "unmarshal status failed")
			continue
		}
		workloads[index].StatusMeta = status
	}
	return workloads, nil
}

func (r *Rediaron) doOpsWorkload(ctx context.Context, workload *types.Workload, processing *types.Processing, create bool) error {
	var err error
	appname, entrypoint, _, err := utils.ParseWorkloadName(workload.Name)
	if err != nil {
		return err
	}

	// now everything is ok
	// we use full length ID instead
	bytes, err := json.Marshal(workload)
	if err != nil {
		return err
	}
	workloadData := string(bytes)

	data := map[string]string{
		fmt.Sprintf(workloadInfoKey, workload.ID):                                                workloadData,
		fmt.Sprintf(nodeWorkloadsKey, workload.Nodename, workload.ID):                            workloadData,
		filepath.Join(workloadDeployPrefix, appname, entrypoint, workload.Nodename, workload.ID): workloadData,
	}

	if create {
		if processing != nil {
			err = r.BatchCreateAndDecr(ctx, data, r.getProcessingKey(processing))
		} else {
			err = r.BatchCreate(ctx, data)
		}
	} else {
		err = r.BatchUpdate(ctx, data)
	}
	return err
}