github.com/m-lab/locate@v0.17.6/cmd/heartbeat/registration/registration.go (about)

     1  package registration
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"fmt"
     7  	"net/url"
     8  	"time"
     9  
    10  	"github.com/google/go-cmp/cmp"
    11  	"github.com/m-lab/go/content"
    12  	"github.com/m-lab/go/host"
    13  	"github.com/m-lab/go/memoryless"
    14  	v2 "github.com/m-lab/locate/api/v2"
    15  	"github.com/m-lab/locate/metrics"
    16  )
    17  
    18  // Loader is a structure to load registration data from siteinfo.
    19  type Loader struct {
    20  	Ticker   *memoryless.Ticker // Ticker determines the interval to reload the data.
    21  	url      *url.URL
    22  	hostname host.Name
    23  	exp      string
    24  	svcs     map[string][]string
    25  	reg      v2.Registration
    26  }
    27  
    28  // NewLoader returns a new loader for registration data.
    29  func NewLoader(ctx context.Context, url *url.URL, hostname, exp string, svcs map[string][]string, config memoryless.Config) (*Loader, error) {
    30  	h, err := host.Parse(hostname)
    31  	if err != nil {
    32  		return nil, err
    33  	}
    34  
    35  	if url == nil {
    36  		return nil, content.ErrUnsupportedURLScheme
    37  	}
    38  
    39  	ticker, err := memoryless.NewTicker(ctx, config)
    40  	if err != nil {
    41  		return nil, err
    42  	}
    43  
    44  	return &Loader{
    45  		Ticker:   ticker,
    46  		url:      url,
    47  		hostname: h,
    48  		exp:      exp,
    49  		svcs:     svcs,
    50  	}, nil
    51  }
    52  
    53  // GetRegistration downloads the registration data from the registration
    54  // URL and matches it with the provided hostname.
    55  func (ldr *Loader) GetRegistration(ctx context.Context) (*v2.Registration, error) {
    56  	provider, err := content.FromURL(ctx, ldr.url)
    57  	if err != nil {
    58  		return nil, err
    59  	}
    60  	exp, err := provider.Get(ctx)
    61  	if err != nil {
    62  		return nil, err
    63  	}
    64  
    65  	var registrations map[string]v2.Registration
    66  	err = json.Unmarshal(exp, &registrations)
    67  	if err != nil {
    68  		return nil, err
    69  	}
    70  
    71  	// The registration key can be both a hostname or a hostname with a service,
    72  	// so the following code checks for both, with priority to hostnames w/o service.
    73  	// Machine name for physical.
    74  	v, ok := registrations[ldr.hostname.String()]
    75  	if !ok {
    76  		// Service name for autonodes.
    77  		v, ok = registrations[ldr.hostname.StringWithService()]
    78  	}
    79  	if ok {
    80  		// Register with fully qualified name.
    81  		v.Hostname = ldr.hostname.StringWithService()
    82  		// If the registration has not changed, there is nothing new to return.
    83  		if cmp.Equal(ldr.reg, v) {
    84  			return nil, nil
    85  		}
    86  
    87  		ldr.reg = v
    88  		v.Experiment = ldr.exp
    89  		v.Services = ldr.svcs
    90  		metrics.RegistrationUpdateTime.Set(float64(time.Now().Unix()))
    91  		return &v, nil
    92  	}
    93  
    94  	return nil, fmt.Errorf("hostname %s not found", ldr.hostname)
    95  }