github.com/grafana/synthetic-monitoring-agent@v0.64.2/internal/pusher/v2/publisher.go (about)

     1  package v2
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  
     7  	"github.com/prometheus/client_golang/prometheus"
     8  	"github.com/rs/zerolog"
     9  
    10  	"github.com/grafana/synthetic-monitoring-agent/internal/model"
    11  	"github.com/grafana/synthetic-monitoring-agent/internal/pusher"
    12  )
    13  
    14  const Name = "v2"
    15  
    16  // NewPublisher creates a new instance of the v2 Publisher.
    17  //
    18  // The provider context is used to control the lifetime of the publisher.
    19  func NewPublisher(ctx context.Context, tenantProvider pusher.TenantProvider, logger zerolog.Logger, pr prometheus.Registerer) pusher.Publisher {
    20  	impl := &publisherImpl{
    21  		ctx:            ctx,
    22  		tenantProvider: tenantProvider,
    23  		options:        defaultPusherOptions,
    24  		handlers:       make(map[model.GlobalID]payloadHandler),
    25  	}
    26  	impl.options.logger = logger
    27  	impl.options.metrics = pusher.NewMetrics(pr)
    28  
    29  	return impl
    30  }
    31  
    32  type payloadHandler interface {
    33  	publish(payload pusher.Payload)
    34  	// run returns the handler that should run or nil to signal that it
    35  	// should be terminated.
    36  	run(ctx context.Context) payloadHandler
    37  }
    38  
    39  type publisherImpl struct {
    40  	ctx            context.Context
    41  	tenantProvider pusher.TenantProvider
    42  	options        pusherOptions
    43  	handlerMutex   sync.Mutex // protects the handlers map
    44  	handlers       map[model.GlobalID]payloadHandler
    45  }
    46  
    47  var _ pusher.Publisher = &publisherImpl{}
    48  
    49  func (p *publisherImpl) Publish(payload pusher.Payload) {
    50  	tenantID := payload.Tenant()
    51  
    52  	handler, found := p.getHandler(tenantID)
    53  	if !found {
    54  		var swapped bool
    55  
    56  		newHandler := newTenantPusher(tenantID, p.tenantProvider, p.options.withTenant(tenantID))
    57  
    58  		handler, swapped = p.replaceHandler(tenantID, nil, newHandler)
    59  		if swapped {
    60  			go p.runHandler(tenantID, handler)
    61  		}
    62  	}
    63  
    64  	handler.publish(payload)
    65  }
    66  
    67  func (p *publisherImpl) runHandler(tenantID model.GlobalID, h payloadHandler) {
    68  	tid, rid := model.GetLocalAndRegionIDs(tenantID)
    69  
    70  	p.options.logger.Info().Int64("tenant_id", tid).Int("region_id", rid).Msg("started push handler")
    71  	defer p.options.logger.Info().Int64("tenant_id", tid).Int("region_id", rid).Msg("stopped push handler")
    72  
    73  	for ok := true; ok && h != nil; {
    74  		next := h.run(p.ctx)
    75  
    76  		h, ok = p.replaceHandler(tenantID, h, next)
    77  		if !ok {
    78  			p.options.logger.Error().Int64("tenant_id", tid).Int("region_id", rid).Msg("unable to swap handler, tenant hijacked")
    79  		}
    80  	}
    81  }
    82  
    83  // replaceHandler exchanges the old handler with the new handler for the tenant
    84  // identified by tenantID.
    85  //
    86  // By passing a nil new handler, you are trying to delete. A nil old handler
    87  // means you are trying to add.
    88  //
    89  // The handler currently in effect is returned, along with whether the handler
    90  // was changed or not.
    91  func (p *publisherImpl) replaceHandler(tenantID model.GlobalID, old, new payloadHandler) (payloadHandler, bool) {
    92  	p.handlerMutex.Lock()
    93  	defer p.handlerMutex.Unlock()
    94  
    95  	// Get the existing handler if any.
    96  	current := p.handlers[tenantID]
    97  
    98  	//  old     | current | new     | op
    99  	//  --------+---------+---------+---------------
   100  	//  nil     | nil     | nil     | delete (noop)
   101  	//  nil     | nil     | non-nil | add
   102  	// *nil     | non-nil | nil     | delete
   103  	// *nil     | non-nil | non-nil | replace
   104  	// *non-nil | nil     | nil     | delete (noop)
   105  	// *non-nil | nil     | non-nil | add
   106  	//  non-nil | non-nil | nil     | delete
   107  	//  non-nil | non-nil | non-nil | replace
   108  
   109  	// If old is nil, that means we are trying to add a handler. If current
   110  	// is not nil, that means there's an existing handler, and the addition
   111  	// is not necessary. If current is nil, we go ahead and add the new handler.
   112  	//
   113  	// If old is not nil, we are trying to replace or delete a handler.
   114  	//
   115  	// If there's nothing there, current is nil and therefore different from
   116  	// old. That means there's nothing to replace (somebody beat us to it)
   117  	// and nothing to delete (idem).
   118  	//
   119  	// If there's something there, current is not nil, we have to replace
   120  	// the existing one.
   121  	if current != old {
   122  		return current, false
   123  	}
   124  
   125  	if new != nil {
   126  		p.handlers[tenantID] = new
   127  	} else {
   128  		delete(p.handlers, tenantID)
   129  	}
   130  
   131  	p.options.metrics.InstalledHandlers.Set(float64(len(p.handlers)))
   132  
   133  	return new, true
   134  }
   135  
   136  func (p *publisherImpl) getHandler(tenantID model.GlobalID) (payloadHandler, bool) {
   137  	p.handlerMutex.Lock()
   138  	defer p.handlerMutex.Unlock()
   139  
   140  	handler, found := p.handlers[tenantID]
   141  
   142  	return handler, found
   143  }