github.com/annwntech/go-micro/v2@v2.9.5/function.go (about)

     1  package micro
     2  
     3  import (
     4  	"context"
     5  	"time"
     6  
     7  	"github.com/annwntech/go-micro/v2/server"
     8  )
     9  
    10  type function struct {
    11  	cancel context.CancelFunc
    12  	Service
    13  }
    14  
    15  func fnHandlerWrapper(f Function) server.HandlerWrapper {
    16  	return func(h server.HandlerFunc) server.HandlerFunc {
    17  		return func(ctx context.Context, req server.Request, rsp interface{}) error {
    18  			defer f.Done()
    19  			return h(ctx, req, rsp)
    20  		}
    21  	}
    22  }
    23  
    24  func fnSubWrapper(f Function) server.SubscriberWrapper {
    25  	return func(s server.SubscriberFunc) server.SubscriberFunc {
    26  		return func(ctx context.Context, msg server.Message) error {
    27  			defer f.Done()
    28  			return s(ctx, msg)
    29  		}
    30  	}
    31  }
    32  
    33  func newFunction(opts ...Option) Function {
    34  	ctx, cancel := context.WithCancel(context.Background())
    35  
    36  	// force ttl/interval
    37  	fopts := []Option{
    38  		RegisterTTL(time.Minute),
    39  		RegisterInterval(time.Second * 30),
    40  	}
    41  
    42  	// prepend to opts
    43  	fopts = append(fopts, opts...)
    44  
    45  	// make context the last thing
    46  	fopts = append(fopts, Context(ctx))
    47  
    48  	service := newService(fopts...)
    49  
    50  	fn := &function{
    51  		cancel:  cancel,
    52  		Service: service,
    53  	}
    54  
    55  	service.Server().Init(
    56  		// ensure the service waits for requests to finish
    57  		server.Wait(nil),
    58  		// wrap handlers and subscribers to finish execution
    59  		server.WrapHandler(fnHandlerWrapper(fn)),
    60  		server.WrapSubscriber(fnSubWrapper(fn)),
    61  	)
    62  
    63  	return fn
    64  }
    65  
    66  func (f *function) Done() error {
    67  	f.cancel()
    68  	return nil
    69  }
    70  
    71  func (f *function) Handle(v interface{}) error {
    72  	return f.Service.Server().Handle(
    73  		f.Service.Server().NewHandler(v),
    74  	)
    75  }
    76  
    77  func (f *function) Subscribe(topic string, v interface{}) error {
    78  	return f.Service.Server().Subscribe(
    79  		f.Service.Server().NewSubscriber(topic, v),
    80  	)
    81  }