github.com/kubeshop/testkube@v1.17.23/pkg/agent/events.go (about)

     1  package agent
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"fmt"
     7  	"time"
     8  
     9  	"github.com/pkg/errors"
    10  	"google.golang.org/grpc"
    11  	"google.golang.org/grpc/encoding/gzip"
    12  	"google.golang.org/grpc/metadata"
    13  
    14  	"github.com/kubeshop/testkube/pkg/api/v1/testkube"
    15  	"github.com/kubeshop/testkube/pkg/cloud"
    16  	"github.com/kubeshop/testkube/pkg/event/kind/common"
    17  )
    18  
    19  var _ common.ListenerLoader = (*Agent)(nil)
    20  
    21  func (ag *Agent) Kind() string {
    22  	return "agent"
    23  }
    24  
    25  func (ag *Agent) Load() (listeners common.Listeners, err error) {
    26  	listeners = append(listeners, ag)
    27  
    28  	return listeners, nil
    29  }
    30  
    31  func (ag *Agent) Name() string {
    32  	return "agent"
    33  }
    34  
    35  func (ag *Agent) Selector() string {
    36  	return ""
    37  }
    38  
    39  func (ag *Agent) Events() []testkube.EventType {
    40  	return testkube.AllEventTypes
    41  }
    42  func (ag *Agent) Metadata() map[string]string {
    43  	return map[string]string{
    44  		"name":     ag.Name(),
    45  		"selector": "",
    46  		"events":   fmt.Sprintf("%v", ag.Events()),
    47  	}
    48  }
    49  
    50  func (ag *Agent) Notify(event testkube.Event) (result testkube.EventResult) {
    51  	event.ClusterName = ag.clusterName
    52  	event.Envs = ag.envs
    53  	// Non blocking send
    54  	select {
    55  	case ag.events <- event:
    56  		return testkube.NewSuccessEventResult(event.Id, "message sent to websocket clients")
    57  	default:
    58  		return testkube.NewFailedEventResult(event.Id, errors.New("message not sent"))
    59  	}
    60  }
    61  
    62  func (ag *Agent) runEventLoop(ctx context.Context) error {
    63  	opts := []grpc.CallOption{grpc.UseCompressor(gzip.Name)}
    64  	md := metadata.Pairs(apiKeyMeta, ag.apiKey)
    65  	ctx = metadata.NewOutgoingContext(ctx, md)
    66  
    67  	stream, err := ag.client.Send(ctx, opts...)
    68  	if err != nil {
    69  		ag.logger.Errorf("failed to execute: %v", err)
    70  		return errors.Wrap(err, "failed to setup stream")
    71  	}
    72  
    73  	ticker := time.NewTicker(30 * time.Second)
    74  	defer ticker.Stop()
    75  
    76  	for {
    77  		select {
    78  		case <-ctx.Done():
    79  			return nil
    80  		case <-ticker.C:
    81  			msg := &cloud.WebsocketData{Opcode: cloud.Opcode_HEALTH_CHECK, Body: nil}
    82  
    83  			err = ag.sendEvent(ctx, stream, msg)
    84  			if err != nil {
    85  				ag.logger.Errorf("websocket stream send healthcheck: %w", err)
    86  
    87  				return err
    88  			}
    89  
    90  		case ev := <-ag.events:
    91  			b, err := json.Marshal(ev)
    92  			if err != nil {
    93  				continue
    94  			}
    95  
    96  			msg := &cloud.WebsocketData{Opcode: cloud.Opcode_TEXT_FRAME, Body: b}
    97  			err = ag.sendEvent(ctx, stream, msg)
    98  			if err != nil {
    99  				ag.logger.Errorf("websocket stream send: %w", err)
   100  				return err
   101  			}
   102  		}
   103  	}
   104  }
   105  
   106  func (ag *Agent) sendEvent(ctx context.Context, stream cloud.TestKubeCloudAPI_SendClient, event *cloud.WebsocketData) error {
   107  	errChan := make(chan error, 1)
   108  	go func() {
   109  		errChan <- stream.Send(event)
   110  		close(errChan)
   111  	}()
   112  
   113  	t := time.NewTimer(ag.sendTimeout)
   114  	select {
   115  	case err := <-errChan:
   116  		if !t.Stop() {
   117  			<-t.C
   118  		}
   119  		return err
   120  	case <-ctx.Done():
   121  		if !t.Stop() {
   122  			<-t.C
   123  		}
   124  
   125  		return ctx.Err()
   126  	case <-t.C:
   127  		return errors.New("too slow")
   128  	}
   129  }