go.dedis.ch/onet/v3@v3.2.11-0.20210930124529-e36530bca7ef/tracing/simul.go (about)

     1  package tracing
     2  
     3  import (
     4  	"bytes"
     5  	"context"
     6  	"encoding/json"
     7  	"fmt"
     8  	"strings"
     9  	"sync"
    10  
    11  	"go.dedis.ch/onet/v3/log"
    12  )
    13  
    14  // The simulCollector is a testing structure that allows to collect all sent
    15  // messages.
    16  // It is mainly used to test the Trace, and *Wrapper structures.
    17  // It simulates a honeycomb integration by simply collecting all fields set,
    18  // and then passing them up when Send is called.
    19  type simulCollector struct {
    20  	Traces    [][]map[string]string
    21  	tracesMux sync.Mutex
    22  	Wg        sync.WaitGroup
    23  }
    24  
    25  func (s *simulCollector) newTrace(ctx context.Context, str string) (context.Context,
    26  	Trace) {
    27  	s.Wg.Add(1)
    28  	newCtx, tr := newSimulTrace(ctx, str)
    29  	go func() {
    30  		<-newCtx.Done()
    31  		s.tracesMux.Lock()
    32  		s.Traces = append(s.Traces, tr.sent)
    33  		s.tracesMux.Unlock()
    34  		s.Wg.Done()
    35  	}()
    36  	return newCtx, tr
    37  }
    38  
    39  func (s *simulCollector) waitAndPrint() {
    40  	s.Wg.Wait()
    41  	for i, c := range s.Traces {
    42  		for j, l := range c {
    43  			fmt.Printf("%d / %d / %s:%s[%s] - %s\n", i, j,
    44  				l["file"], l["method"], l["line"], l["log.Msg"])
    45  			//fmt.Printf("%d/%d %+v\n", i, j, l)
    46  		}
    47  	}
    48  }
    49  
    50  type simulTrace struct {
    51  	ctx    context.Context
    52  	cancel context.CancelFunc
    53  	kvs    map[string]string
    54  	root   *simulSpan
    55  	sent   []map[string]string
    56  }
    57  
    58  func newSimulTrace(ctx context.Context, s string) (context.Context, *simulTrace) {
    59  	newCtx, cancel := context.WithCancel(ctx)
    60  	tt := &simulTrace{
    61  		kvs:    make(map[string]string),
    62  		root:   newSimulSpan(),
    63  		ctx:    newCtx,
    64  		cancel: cancel,
    65  	}
    66  	tt.root.trace = tt
    67  	return newCtx, tt
    68  }
    69  
    70  func (tt *simulTrace) AddField(key string, val interface{}) {
    71  	tt.kvs[key] = valToJSON(val)
    72  }
    73  
    74  func (tt *simulTrace) GetRootSpan() Span {
    75  	return tt.root
    76  }
    77  func (tt *simulTrace) Send() {
    78  	tt.root.Send()
    79  	for i := range tt.sent {
    80  		for k, v := range tt.kvs {
    81  			tt.sent[i][k] = v
    82  		}
    83  	}
    84  	tt.cancel()
    85  }
    86  
    87  func (tt *simulTrace) printTrace() {
    88  	tt.root.printTrace()
    89  }
    90  
    91  type simulSpan struct {
    92  	kvs    map[string]string
    93  	child  *simulSpan
    94  	parent *simulSpan
    95  	trace  *simulTrace
    96  }
    97  
    98  func newSimulSpan() *simulSpan {
    99  	return &simulSpan{
   100  		kvs: make(map[string]string),
   101  	}
   102  }
   103  
   104  func (ts *simulSpan) AddField(key string, val interface{}) {
   105  	ts.kvs[key] = valToJSON(val)
   106  }
   107  func (ts *simulSpan) CreateChild(ctx context.Context) (context.Context, Span) {
   108  	ts.child = newSimulSpan()
   109  	ts.child.parent = ts
   110  	ts.child.trace = ts.trace
   111  	return ctx, ts.child
   112  }
   113  func (ts *simulSpan) Send() {
   114  	if ts.child != nil {
   115  		ts.child.Send()
   116  	}
   117  	ts.trace.sent = append(ts.trace.sent, ts.kvs)
   118  	if ts.parent != nil {
   119  		ts.parent.child = nil
   120  	}
   121  }
   122  
   123  func (ts *simulSpan) printTrace() {
   124  	if ts.child != nil {
   125  		ts.child.printTrace()
   126  	}
   127  }
   128  
   129  func valToJSON(val interface{}) string {
   130  	var buf = &bytes.Buffer{}
   131  	var enc = json.NewEncoder(buf)
   132  	log.ErrFatal(enc.Encode(val))
   133  	return strings.TrimSpace(buf.String())
   134  }
   135  
   136  // newSimulLogger returns a new simulationCollector and a TraceLogger.
   137  // You can wait on the simulCollector.
   138  // Wg and get all traces in the simulCollector.Traces
   139  func newSimulLogger() (*simulCollector, *TraceLogger) {
   140  	sc := &simulCollector{}
   141  	return sc, NewLogger(sc.newTrace)
   142  }