github.com/netdata/go.d.plugin@v0.58.1/agent/functions/manager.go (about)

     1  // SPDX-License-Identifier: GPL-3.0-or-later
     2  
     3  package functions
     4  
     5  import (
     6  	"bufio"
     7  	"context"
     8  	"io"
     9  	"log/slog"
    10  	"os"
    11  	"strings"
    12  	"sync"
    13  
    14  	"github.com/netdata/go.d.plugin/logger"
    15  
    16  	"github.com/mattn/go-isatty"
    17  	"github.com/muesli/cancelreader"
    18  )
    19  
    20  var isTerminal = isatty.IsTerminal(os.Stdout.Fd()) || isatty.IsTerminal(os.Stdin.Fd())
    21  
    22  func NewManager() *Manager {
    23  	return &Manager{
    24  		Logger: logger.New().With(
    25  			slog.String("component", "functions manager"),
    26  		),
    27  		Input:            os.Stdin,
    28  		mux:              &sync.Mutex{},
    29  		FunctionRegistry: make(map[string]func(Function)),
    30  	}
    31  }
    32  
    33  type Manager struct {
    34  	*logger.Logger
    35  
    36  	Input            io.Reader
    37  	mux              *sync.Mutex
    38  	FunctionRegistry map[string]func(Function)
    39  }
    40  
    41  func (m *Manager) Register(name string, fn func(Function)) {
    42  	if fn == nil {
    43  		m.Warningf("not registering '%s': nil function", name)
    44  		return
    45  	}
    46  	m.addFunction(name, fn)
    47  }
    48  
    49  func (m *Manager) Run(ctx context.Context) {
    50  	m.Info("instance is started")
    51  	defer func() { m.Info("instance is stopped") }()
    52  
    53  	if !isTerminal {
    54  		var wg sync.WaitGroup
    55  
    56  		r, err := cancelreader.NewReader(m.Input)
    57  		if err != nil {
    58  			m.Errorf("fail to create cancel reader: %v", err)
    59  			return
    60  		}
    61  
    62  		go func() { <-ctx.Done(); r.Cancel() }()
    63  
    64  		wg.Add(1)
    65  		go func() { defer wg.Done(); m.run(r) }()
    66  
    67  		wg.Wait()
    68  		_ = r.Close()
    69  	}
    70  
    71  	<-ctx.Done()
    72  }
    73  
    74  func (m *Manager) run(r io.Reader) {
    75  	sc := bufio.NewScanner(r)
    76  
    77  	for sc.Scan() {
    78  		text := sc.Text()
    79  
    80  		var fn *Function
    81  		var err error
    82  
    83  		// FIXME:  if we are waiting for FUNCTION_PAYLOAD_END and a new FUNCTION* appears,
    84  		// we need to discard the current one and switch to the new one
    85  		switch {
    86  		case strings.HasPrefix(text, "FUNCTION "):
    87  			fn, err = parseFunction(text)
    88  		case strings.HasPrefix(text, "FUNCTION_PAYLOAD "):
    89  			fn, err = parseFunctionWithPayload(text, sc)
    90  		case text == "":
    91  			continue
    92  		default:
    93  			m.Warningf("unexpected line: '%s'", text)
    94  			continue
    95  		}
    96  
    97  		if err != nil {
    98  			m.Warningf("parse function: %v ('%s')", err, text)
    99  			continue
   100  		}
   101  
   102  		function, ok := m.lookupFunction(fn.Name)
   103  		if !ok {
   104  			m.Infof("skipping execution of '%s': unregistered function", fn.Name)
   105  			continue
   106  		}
   107  		if function == nil {
   108  			m.Warningf("skipping execution of '%s': nil function registered", fn.Name)
   109  			continue
   110  		}
   111  
   112  		m.Debugf("executing function: '%s'", fn.String())
   113  		function(*fn)
   114  	}
   115  }
   116  
   117  func (m *Manager) addFunction(name string, fn func(Function)) {
   118  	m.mux.Lock()
   119  	defer m.mux.Unlock()
   120  
   121  	if _, ok := m.FunctionRegistry[name]; !ok {
   122  		m.Debugf("registering function '%s'", name)
   123  	} else {
   124  		m.Warningf("re-registering function '%s'", name)
   125  	}
   126  	m.FunctionRegistry[name] = fn
   127  }
   128  
   129  func (m *Manager) lookupFunction(name string) (func(Function), bool) {
   130  	m.mux.Lock()
   131  	defer m.mux.Unlock()
   132  
   133  	f, ok := m.FunctionRegistry[name]
   134  	return f, ok
   135  }