argo/client/watch_workflow.go

package client

import (
	"context"
	"fmt"
	"github.com/argoproj/argo-workflows/v3/pkg/apiclient/workflow"
	"github.com/argoproj/argo-workflows/v3/pkg/apis/workflow/v1alpha1"
	"github.com/argoproj/argo-workflows/v3/util"
	"github.com/pkg/errors"
	"io"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"time"
)

type WatchWorkflowCallback func(wf *v1alpha1.Workflow)

type WatchWorkflowToken struct {
	workflowChan        chan *v1alpha1.Workflow
	doneChan            chan any
	streamReceiveCancel context.CancelFunc
}

func (token *WatchWorkflowToken) Done() {
	token.streamReceiveCancel()

	token.doneChan <- nil
	close(token.doneChan)
	token.doneChan = nil

	close(token.workflowChan)
	token.workflowChan = nil
}

type WatchWorkflowParams struct {
	Namespace string
	Name      string
}

// WatchWorkflow 监听工作流
func (c *Client) WatchWorkflow(params WatchWorkflowParams, callback WatchWorkflowCallback) (*WatchWorkflowToken, error) {
	req := &workflow.WatchWorkflowsRequest{
		Namespace: params.Namespace,
		ListOptions: &metav1.ListOptions{
			FieldSelector:   util.GenerateFieldSelectorFromWorkflowName(params.Name),
			ResourceVersion: "0",
		},
	}

	cancelCtx, cancel := context.WithCancel(c.ctx)
	stream, err := c.workflowService.WatchWorkflows(cancelCtx, req)
	if err != nil {
		cancel()
		return nil, errors.New(err.Error())
	}

	token := &WatchWorkflowToken{
		workflowChan:        make(chan *v1alpha1.Workflow),
		doneChan:            make(chan any),
		streamReceiveCancel: cancel,
	}

	go func() {
		for {
			select {
			case <-token.doneChan:
				return
			case wf := <-token.workflowChan:
				if wf == nil {
					continue
				}

				if callback != nil {
					callback(wf)
				}
			}
		}
	}()

	go func() {
		for {
			event, err := stream.Recv()
			if err != nil {
				if err == io.EOF {
					if stream.Context().Err().Error() == "context canceled" {
						return
					}

					cancelCtx, cancel := context.WithCancel(c.ctx)
					stream, err = c.workflowService.WatchWorkflows(cancelCtx, req)
					if err != nil {
						cancel()
						fmt.Printf("%v\n", errors.New(err.Error()))
						time.Sleep(5 * time.Second)
						continue
					}

					token.streamReceiveCancel = cancel

					continue
				}

				fmt.Printf("%v\n", errors.New(err.Error()))
				return
			}

			if event == nil {
				continue
			}

			token.workflowChan <- event.Object
		}
	}()

	return token, nil
}