github.com/aacfactory/fns@v1.2.86-0.20240310083819-80d667fc0a17/services/tracings/middleware.go (about)

     1  /*
     2   * Copyright 2023 Wang Min Xiang
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   * 	http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   *
    16   */
    17  
    18  package tracings
    19  
    20  import (
    21  	"github.com/aacfactory/configures"
    22  	"github.com/aacfactory/errors"
    23  	"github.com/aacfactory/fns/context"
    24  	"github.com/aacfactory/fns/transports"
    25  	"github.com/aacfactory/json"
    26  	"github.com/aacfactory/logs"
    27  )
    28  
    29  type Config struct {
    30  	Enable      bool            `json:"enable"`
    31  	BatchSize   int             `json:"batchSize"`
    32  	ChannelSize int             `json:"channelSize"`
    33  	Reporter    json.RawMessage `json:"reporter"`
    34  }
    35  
    36  type Middleware struct {
    37  	log      logs.Logger
    38  	enable   bool
    39  	events   chan *Trace
    40  	cancel   context.CancelFunc
    41  	reporter Reporter
    42  }
    43  
    44  func (middle *Middleware) Name() string {
    45  	return "tracings"
    46  }
    47  
    48  func (middle *Middleware) Construct(options transports.MiddlewareOptions) (err error) {
    49  	middle.log = options.Log
    50  	config := Config{}
    51  	configErr := options.Config.As(&config)
    52  	if configErr != nil {
    53  		err = errors.Warning("fns: tracing middleware construct failed").WithCause(configErr)
    54  		return
    55  	}
    56  	if config.Enable {
    57  		reporterConfig, reporterConfigErr := configures.NewJsonConfig(config.Reporter)
    58  		if configErr != nil {
    59  			err = errors.Warning("fns: tracing middleware construct failed").WithCause(reporterConfigErr)
    60  			return
    61  		}
    62  		reportErr := middle.reporter.Construct(ReporterOptions{
    63  			Log:    middle.log,
    64  			Config: reporterConfig,
    65  		})
    66  		if reportErr != nil {
    67  			err = errors.Warning("fns: tracing middleware construct failed").WithCause(reportErr)
    68  			return
    69  		}
    70  		batchSize := config.BatchSize
    71  		if batchSize < 0 {
    72  			batchSize = 4
    73  		}
    74  		chs := config.ChannelSize
    75  		if chs < 0 {
    76  			chs = 4096
    77  		}
    78  		middle.events = make(chan *Trace, chs)
    79  		ctx, cancel := context.WithCancel(context.TODO())
    80  		middle.cancel = cancel
    81  		for i := 0; i < batchSize; i++ {
    82  			middle.listen(ctx)
    83  		}
    84  	}
    85  	return
    86  }
    87  
    88  func (middle *Middleware) Handler(next transports.Handler) transports.Handler {
    89  	if middle.enable {
    90  		return transports.HandlerFunc(func(w transports.ResponseWriter, r transports.Request) {
    91  			id := r.Header().Get(transports.RequestIdHeaderName)
    92  			if len(id) == 0 {
    93  				next.Handle(w, r)
    94  				return
    95  			}
    96  			tracer := New(id)
    97  			With(r, tracer)
    98  			next.Handle(w, r)
    99  			trace := tracer.Trace()
   100  			middle.events <- trace
   101  		})
   102  	}
   103  	return next
   104  }
   105  
   106  func (middle *Middleware) Close() {
   107  	if middle.cancel != nil {
   108  		middle.cancel()
   109  	}
   110  }
   111  
   112  func (middle *Middleware) listen(ctx context.Context) {
   113  	go func(ctx context.Context, events chan *Trace, reporter Reporter) {
   114  		stop := false
   115  		for {
   116  			select {
   117  			case <-ctx.Done():
   118  				stop = true
   119  				break
   120  			case trace, ok := <-events:
   121  				if !ok {
   122  					stop = true
   123  					break
   124  				}
   125  				reporter.Report(ctx, trace)
   126  				break
   127  			}
   128  			if stop {
   129  				break
   130  			}
   131  		}
   132  	}(ctx, middle.events, middle.reporter)
   133  }