github.com/blend/go-sdk@v1.20220411.3/expvar/vars.go (about)

     1  /*
     2  
     3  Copyright (c) 2022 - Present. Blend Labs, Inc. All rights reserved
     4  Use of this source code is governed by a MIT license that can be found in the LICENSE file.
     5  
     6  */
     7  
     8  package expvar
     9  
    10  import (
    11  	"fmt"
    12  	"io"
    13  	"net/http"
    14  	"sort"
    15  	"sync"
    16  )
    17  
    18  // Vars is a collection of expvars.
    19  type Vars struct {
    20  	vars      sync.Map // map[string]Var
    21  	varKeysMu sync.RWMutex
    22  	varKeys   []string // sorted
    23  }
    24  
    25  // Get retrieves a named exported variable. It returns nil if the name has
    26  // not been registered.
    27  func (v *Vars) Get(name string) Var {
    28  	i, _ := v.vars.Load(name)
    29  	val, _ := i.(Var)
    30  	return val
    31  }
    32  
    33  // Publish declares a named exported variable. This should be called from a
    34  // package's init function when it creates its Vars. If the name is already
    35  // registered then this will log.Panic.
    36  func (v *Vars) Publish(name string, val Var) error {
    37  	if _, dup := v.vars.LoadOrStore(name, val); dup {
    38  		return fmt.Errorf("reuse of exported var name: %s", name)
    39  	}
    40  	v.varKeysMu.Lock()
    41  	defer v.varKeysMu.Unlock()
    42  	v.varKeys = append(v.varKeys, name)
    43  	sort.Strings(v.varKeys)
    44  	return nil
    45  }
    46  
    47  // Forward forwards the vars contained in this set to another set with a given key prefix.
    48  func (v *Vars) Forward(dst *Vars, keyPrefix string) error {
    49  	var err error
    50  	return v.Do(func(kv KeyValue) error {
    51  		if err = dst.Publish(keyPrefix+kv.Key, kv.Value); err != nil {
    52  			return err
    53  		}
    54  		return nil
    55  	})
    56  }
    57  
    58  // Do calls f for each exported variable.
    59  // The global variable map is locked during the iteration,
    60  // but existing entries may be concurrently updated.
    61  func (v *Vars) Do(f func(KeyValue) error) error {
    62  	v.varKeysMu.RLock()
    63  	defer v.varKeysMu.RUnlock()
    64  	var err error
    65  	for _, k := range v.varKeys {
    66  		val, _ := v.vars.Load(k)
    67  		err = f(KeyValue{k, val.(Var)})
    68  		if err != nil {
    69  			return err
    70  		}
    71  	}
    72  	return nil
    73  }
    74  
    75  // Handler returns an http.HandlerFunc that renders the vars as json.
    76  func (v *Vars) Handler(w http.ResponseWriter, r *http.Request) {
    77  	w.Header().Set("Content-Type", "application/json; charset=utf-8")
    78  	_, _ = v.WriteTo(w)
    79  }
    80  
    81  // WriteTo writes the vars to a given writer as json.
    82  //
    83  // This is called by the Handler function to return the output.
    84  func (v *Vars) WriteTo(wr io.Writer) (size int64, err error) {
    85  	var n int
    86  	n, err = fmt.Fprint(wr, "{")
    87  	size += int64(n)
    88  	if err != nil {
    89  		return
    90  	}
    91  	first := true
    92  	err = v.Do(func(kv KeyValue) error {
    93  		if !first {
    94  			n, err = fmt.Fprint(wr, ",")
    95  			size += int64(n)
    96  			if err != nil {
    97  				return err
    98  			}
    99  		}
   100  		first = false
   101  		n, err = fmt.Fprintf(wr, "%q:%s", kv.Key, kv.Value)
   102  		size += int64(n)
   103  		if err != nil {
   104  			return err
   105  		}
   106  		return nil
   107  	})
   108  	if err != nil {
   109  		return
   110  	}
   111  	n, err = fmt.Fprintln(wr, "}")
   112  	size += int64(n)
   113  	return
   114  }