pkg/capabilities/consensus/ocr3: use services.Engine (#1009)

Author: jmank88

pkg/capabilities/consensus/ocr3/capability.go

@@ -33,14 +33,13 @@ var info = capabilities.MustNewCapabilityInfo(
 )
 
 type capability struct {
-	services.StateMachine
+	services.Service
+	eng *services.Engine
+
 	capabilities.CapabilityInfo
 	capabilities.Validator[config, inputs, requests.Response]
 
 	reqHandler *requests.Handler
-	stopCh     services.StopChan
-	wg         sync.WaitGroup
-	lggr       logger.Logger
 
 	requestTimeout     time.Duration
 	requestTimeoutLock sync.RWMutex
@@ -67,11 +66,8 @@ func NewCapability(s *requests.Store, clock clockwork.Clock, requestTimeout time
 	o := &capability{
 		CapabilityInfo:    info,
 		Validator:         capabilities.NewValidator[config, inputs, requests.Response](capabilities.ValidatorArgs{Info: info}),
-		reqHandler:        requests.NewHandler(lggr, s, clock, requestTimeout),
 		clock:             clock,
 		requestTimeout:    requestTimeout,
-		stopCh:            make(chan struct{}),
-		lggr:              logger.Named(lggr, "OCR3CapabilityClient"),
 		aggregatorFactory: aggregatorFactory,
 		aggregators:       map[string]types.Aggregator{},
 		encoderFactory:    encoderFactory,
@@ -80,39 +76,16 @@ func NewCapability(s *requests.Store, clock clockwork.Clock, requestTimeout time
 		callbackChannelBufferSize: callbackChannelBufferSize,
 		registeredWorkflowsIDs:    map[string]bool{},
 	}
+	o.Service, o.eng = services.Config{
+		Name: "OCR3CapabilityClient",
+		NewSubServices: func(l logger.Logger) []services.Service {
+			o.reqHandler = requests.NewHandler(lggr, s, clock, requestTimeout)
+			return []services.Service{o.reqHandler}
+		},
+	}.NewServiceEngine(lggr)
 	return o
 }
 
-func (o *capability) Start(ctx context.Context) error {
-	return o.StartOnce("OCR3Capability", func() error {
-		err := o.reqHandler.Start(ctx)
-		if err != nil {
-			return fmt.Errorf("failed to start request handler: %w", err)
-		}
-
-		return nil
-	})
-}
-
-func (o *capability) Close() error {
-	return o.StopOnce("OCR3Capability", func() error {
-		close(o.stopCh)
-		o.wg.Wait()
-		err := o.reqHandler.Close()
-		if err != nil {
-			return fmt.Errorf("failed to close request handler: %w", err)
-		}
-
-		return nil
-	})
-}
-
-func (o *capability) Name() string { return o.lggr.Name() }
-
-func (o *capability) HealthReport() map[string]error {
-	return map[string]error{o.Name(): o.Healthy()}
-}
-
 func (o *capability) RegisterToWorkflow(ctx context.Context, request capabilities.RegisterToWorkflowRequest) error {
 	c, err := o.ValidateConfig(request.Config)
 	if err != nil {
@@ -121,13 +94,13 @@ func (o *capability) RegisterToWorkflow(ctx context.Context, request capabilitie
 
 	o.mu.Lock()
 	defer o.mu.Unlock()
-	agg, err := o.aggregatorFactory(c.AggregationMethod, *c.AggregationConfig, o.lggr)
+	agg, err := o.aggregatorFactory(c.AggregationMethod, *c.AggregationConfig, o.eng)
 	if err != nil {
 		return err
 	}
 	o.aggregators[request.Metadata.WorkflowID] = agg
 
-	encoder, err := o.encoderFactory(c.Encoder, c.EncoderConfig, o.lggr)
+	encoder, err := o.encoderFactory(c.Encoder, c.EncoderConfig, o.eng)
 	if err != nil {
 		return err
 	}
@@ -155,7 +128,7 @@ func (o *capability) GetEncoderByWorkflowID(workflowID string) (types.Encoder, e
 }
 
 func (o *capability) GetEncoderByName(encoderName string, config *values.Map) (types.Encoder, error) {
-	return o.encoderFactory(encoderName, config, o.lggr)
+	return o.encoderFactory(encoderName, config, o.eng)
 }
 
 func (o *capability) GetRegisteredWorkflowsIDs() []string {
@@ -205,7 +178,7 @@ func (o *capability) Execute(ctx context.Context, r capabilities.CapabilityReque
 	}
 	err := r.Inputs.UnwrapTo(&m)
 	if err != nil {
-		o.lggr.Warnf("could not unwrap method from CapabilityRequest, using default: %v", err)
+		o.eng.Warnf("could not unwrap method from CapabilityRequest, using default: %v", err)
 	}
 
 	switch m.Method {
@@ -214,10 +187,10 @@ func (o *capability) Execute(ctx context.Context, r capabilities.CapabilityReque
 		if err != nil {
 			return capabilities.CapabilityResponse{}, fmt.Errorf("failed to create map for response inputs: %w", err)
 		}
-		o.lggr.Debugw("Execute - sending response", "workflowExecutionID", r.Metadata.WorkflowExecutionID, "inputs", inputs, "terminate", m.Terminate)
+		o.eng.Debugw("Execute - sending response", "workflowExecutionID", r.Metadata.WorkflowExecutionID, "inputs", inputs, "terminate", m.Terminate)
 		var responseErr error
 		if m.Terminate {
-			o.lggr.Debugw("Execute - terminating execution", "workflowExecutionID", r.Metadata.WorkflowExecutionID)
+			o.eng.Debugw("Execute - terminating execution", "workflowExecutionID", r.Metadata.WorkflowExecutionID)
 			responseErr = capabilities.ErrStopExecution
 		}
 		out := requests.Response{
@@ -298,7 +271,7 @@ func (o *capability) queueRequestForProcessing(
 		ExpiresAt:                o.clock.Now().Add(requestTimeout),
 	}
 
-	o.lggr.Debugw("Execute - adding to store", "workflowID", r.WorkflowID, "workflowExecutionID", r.WorkflowExecutionID, "observations", r.Observations)
+	o.eng.Debugw("Execute - adding to store", "workflowID", r.WorkflowID, "workflowExecutionID", r.WorkflowExecutionID, "observations", r.Observations)
 
 	o.reqHandler.SendRequest(ctx, r)
 	return callbackCh, nil

pkg/capabilities/consensus/ocr3/requests/handler.go

@@ -3,7 +3,6 @@ package requests
 import (
 	"context"
 	"fmt"
-	"sync"
 	"time"
 
 	"github.com/jonboulle/clockwork"
@@ -18,9 +17,8 @@ type responseCacheEntry struct {
 }
 
 type Handler struct {
-	services.StateMachine
-
-	lggr logger.Logger
+	services.Service
+	eng *services.Engine
 
 	store *Store
 
@@ -33,23 +31,23 @@ type Handler struct {
 	requestCh  chan *Request
 
 	clock clockwork.Clock
-
-	stopCh services.StopChan
-	wg     sync.WaitGroup
 }
 
 func NewHandler(lggr logger.Logger, s *Store, clock clockwork.Clock, responseExpiryTime time.Duration) *Handler {
-	return &Handler{
-		lggr:            lggr,
+	h := &Handler{
 		store:           s,
 		pendingRequests: map[string]*Request{},
 		responseCache:   map[string]*responseCacheEntry{},
 		responseCh:      make(chan Response),
 		requestCh:       make(chan *Request),
 		clock:           clock,
 		cacheExpiryTime: responseExpiryTime,
-		stopCh:          make(services.StopChan),
 	}
+	h.Service, h.eng = services.Config{
+		Name:  "Handler",
+		Start: h.start,
+	}.NewServiceEngine(lggr)
+	return h
 }
 
 func (h *Handler) SendResponse(ctx context.Context, resp Response) {
@@ -68,25 +66,9 @@ func (h *Handler) SendRequest(ctx context.Context, r *Request) {
 	}
 }
 
-func (h *Handler) Start(_ context.Context) error {
-	return h.StartOnce("RequestHandler", func() error {
-		h.wg.Add(1)
-		go func() {
-			defer h.wg.Done()
-			ctx, cancel := h.stopCh.NewCtx()
-			defer cancel()
-			h.worker(ctx)
-		}()
-		return nil
-	})
-}
-
-func (h *Handler) Close() error {
-	return h.StopOnce("RequestHandler", func() error {
-		close(h.stopCh)
-		h.wg.Wait()
-		return nil
-	})
+func (h *Handler) start(_ context.Context) error {
+	h.eng.Go(h.worker)
+	return nil
 }
 
 func (h *Handler) worker(ctx context.Context) {
@@ -111,13 +93,13 @@ func (h *Handler) worker(ctx context.Context) {
 			existingResponse := h.responseCache[req.WorkflowExecutionID]
 			if existingResponse != nil {
 				delete(h.responseCache, req.WorkflowExecutionID)
-				h.lggr.Debugw("Found cached response for request", "workflowExecutionID", req.WorkflowExecutionID)
+				h.eng.Debugw("Found cached response for request", "workflowExecutionID", req.WorkflowExecutionID)
 				h.sendResponse(ctx, req, existingResponse.response)
 				continue
 			}
 
 			if err := h.store.Add(req); err != nil {
-				h.lggr.Errorw("failed to add request to store", "err", err)
+				h.eng.Errorw("failed to add request to store", "err", err)
 			}
 
 		case resp := <-h.responseCh:
@@ -127,7 +109,7 @@ func (h *Handler) worker(ctx context.Context) {
 					response:  resp,
 					entryTime: h.clock.Now(),
 				}
-				h.lggr.Debugw("Caching response without request", "workflowExecutionID", resp.WorkflowExecutionID)
+				h.eng.Debugw("Caching response without request", "workflowExecutionID", resp.WorkflowExecutionID)
 				continue
 			}
 
@@ -165,7 +147,7 @@ func (h *Handler) expireCachedResponses() {
 	for k, v := range h.responseCache {
 		if h.clock.Since(v.entryTime) > h.cacheExpiryTime {
 			delete(h.responseCache, k)
-			h.lggr.Debugw("Expired response", "workflowExecutionID", k)
+			h.eng.Debugw("Expired response", "workflowExecutionID", k)
 		}
 	}
 }