github.com/insolar/vanilla@v0.0.0-20201023172447-248fdf805322/watchdog/context.go (about)

     1  // Copyright 2020 Insolar Network Ltd.
     2  // All rights reserved.
     3  // This material is licensed under the Insolar License version 1.0,
     4  // available at https://github.com/insolar/assured-ledger/blob/master/LICENSE.md.
     5  
     6  package watchdog
     7  
     8  import (
     9  	"context"
    10  	"time"
    11  )
    12  
    13  func DoneOf(ctx context.Context) <-chan struct{} {
    14  	if h := from(ctx); h != nil {
    15  		return h.done(ctx)
    16  	}
    17  	return ctx.Done()
    18  }
    19  
    20  func Beat(ctx context.Context) {
    21  	ForcedBeat(ctx, false)
    22  }
    23  
    24  func ForcedBeat(ctx context.Context, forced bool) {
    25  	if h := from(ctx); h != nil {
    26  		h.beat(ctx, forced)
    27  	}
    28  }
    29  
    30  func WithFrame(ctx context.Context, frameName string) context.Context {
    31  	if h := from(ctx); h != nil {
    32  		return h.root.createSubFrame(ctx, frameName, h)
    33  	}
    34  	return ctx
    35  }
    36  
    37  func Call(ctx context.Context, frameName string, fn func(context.Context)) {
    38  	if frame := from(ctx); frame == nil {
    39  		fn(ctx)
    40  	} else {
    41  		frame.root.createSubFrame(ctx, frameName, frame).call(fn)
    42  	}
    43  }
    44  
    45  func WithFactory(ctx context.Context, name string, factory HeartbeatGeneratorFactory) context.Context {
    46  	r := frameRoot{factory}
    47  	return r.createSubFrame(ctx, name, nil)
    48  }
    49  
    50  func WithoutFactory(ctx context.Context) context.Context {
    51  	return context.WithValue(ctx, watchdogKey, nil) // stop search
    52  }
    53  
    54  func FromContext(ctx context.Context) (bool, HeartbeatGeneratorFactory) {
    55  	switch f := from(ctx); {
    56  	case f == nil:
    57  		return false, nil
    58  	case f.root == nil:
    59  		return true, nil
    60  	default:
    61  		return true, f.root.factory
    62  	}
    63  }
    64  
    65  func from(ctx context.Context) *frame {
    66  	if h, ok := ctx.Value(watchdogKey).(*frame); ok {
    67  		return h
    68  	}
    69  	return nil
    70  }
    71  
    72  var watchdogKey = &struct{}{}
    73  
    74  type frameRoot struct {
    75  	factory HeartbeatGeneratorFactory
    76  }
    77  
    78  func (r *frameRoot) createSubFrame(ctx context.Context, name string, parent *frame) *frame {
    79  	if parent != nil {
    80  		name = parent.name + "/" + name
    81  	}
    82  	return &frame{r, ctx, r.factory.CreateGenerator(name), name}
    83  }
    84  
    85  type frame struct {
    86  	root      *frameRoot
    87  	context   context.Context
    88  	generator *HeartbeatGenerator
    89  	name      string
    90  }
    91  
    92  func (h *frame) Deadline() (deadline time.Time, ok bool) {
    93  	h.beat(h.context, false)
    94  	return h.context.Deadline()
    95  }
    96  
    97  func (h *frame) Value(key interface{}) interface{} {
    98  	if watchdogKey == key {
    99  		return h
   100  	}
   101  	h.beat(h.context, false)
   102  	return h.context.Value(key)
   103  }
   104  
   105  func (h *frame) Err() error {
   106  	err := h.context.Err()
   107  	if err != nil {
   108  		h.generator.Cancel()
   109  	} else {
   110  		h.generator.Heartbeat()
   111  	}
   112  	return err
   113  }
   114  
   115  func (h *frame) Done() <-chan struct{} {
   116  	return h.done(h.context)
   117  }
   118  
   119  func (h *frame) beat(ctx context.Context, forced bool) {
   120  	if ctx.Err() != nil {
   121  		h.generator.Cancel()
   122  	} else {
   123  		h.generator.ForcedHeartbeat(forced)
   124  	}
   125  }
   126  
   127  func (h *frame) start() {
   128  	h.generator.ForcedHeartbeat(true)
   129  }
   130  
   131  func (h *frame) cancel() {
   132  	h.generator.Cancel()
   133  }
   134  
   135  func (h *frame) done(ctx context.Context) <-chan struct{} {
   136  	ch := ctx.Done()
   137  	select {
   138  	case <-ch:
   139  		h.generator.Cancel()
   140  	default:
   141  		h.generator.ForcedHeartbeat(false)
   142  	}
   143  	return ch
   144  }
   145  
   146  func (h *frame) call(fn func(context.Context)) {
   147  	h.start()
   148  	defer h.cancel()
   149  	fn(h)
   150  }