github.com/micro/go-micro/examples@v0.0.0-20210105173217-bf4ab679e18b/waitgroup/main.go (about)

     1  package main
     2  
     3  import (
     4  	"fmt"
     5  	"sync"
     6  
     7  	"context"
     8  	"github.com/micro/go-micro/v2"
     9  	"github.com/micro/go-micro/v2/server"
    10  )
    11  
    12  // waitgroup is a handler wrapper which adds a handler to a sync.WaitGroup
    13  func waitgroup(wg *sync.WaitGroup) server.HandlerWrapper {
    14  	return func(h server.HandlerFunc) server.HandlerFunc {
    15  		return func(ctx context.Context, req server.Request, rsp interface{}) error {
    16  			wg.Add(1)
    17  			defer wg.Done()
    18  			return h(ctx, req, rsp)
    19  		}
    20  	}
    21  }
    22  
    23  func main() {
    24  	var wg sync.WaitGroup
    25  
    26  	service := micro.NewService(
    27  		// wrap handlers with waitgroup wrapper
    28  		micro.WrapHandler(waitgroup(&wg)),
    29  		// waits for the waitgroup once stopped
    30  		micro.AfterStop(func() error {
    31  			// wait for handlers to finish
    32  			wg.Wait()
    33  			return nil
    34  		}),
    35  	)
    36  
    37  	if err := service.Run(); err != nil {
    38  		fmt.Println(err)
    39  	}
    40  }