add telemetry to nats,http and lua executor

Telemetry needs the TELEMETRY_TRACES env variable set to anything in
order to enable the traces. By default it will output to stdout but
traces can be sent to an OpenTelemetry server (like Grafana's Tempo).

Added those options to the NixOS module

Author: numkem

cmd/server/http.go

@@ -12,13 +12,21 @@ import (
 	log "github.com/sirupsen/logrus"
 	"golang.org/x/net/context"
 
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/codes"
+	"go.opentelemetry.io/otel/propagation"
+	"go.opentelemetry.io/otel/trace"
+
 	"github.com/numkem/msgscript"
 	"github.com/numkem/msgscript/executor"
 )
 
 const DEFAULT_HTTP_PORT = 7643
 const DEFAULT_HTTP_TIMEOUT = 5 * time.Second
 
+var tracer = otel.Tracer("http-nats-proxy")
+
 type httpNatsProxy struct {
 	port string
 	nc   *nats.Conn
@@ -40,34 +48,56 @@ func NewHttpNatsProxy(port int, natsURL string) (*httpNatsProxy, error) {
 }
 
 func (p *httpNatsProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
+	// Extract context from incoming request headers
+	ctx := otel.GetTextMapPropagator().Extract(r.Context(), propagation.HeaderCarrier(r.Header))
+	
+	// Start a new span for the HTTP request
+	ctx, span := tracer.Start(ctx, "http.request",
+		trace.WithSpanKind(trace.SpanKindServer),
+		trace.WithAttributes(
+			attribute.String("http.method", r.Method),
+			attribute.String("http.url", r.URL.String()),
+			attribute.String("http.remote_addr", r.RemoteAddr),
+		),
+	)
+	defer span.End()
+
 	defer r.Body.Close()
 
 	// URL should look like /funcs.foobar
 	// Where funcs.foobar is the subject for NATS
 	ss := strings.Split(r.URL.Path, "/")
 	// Validate URL structure
 	if len(ss) < 2 {
+		span.SetStatus(codes.Error, "Invalid URL structure")
+		span.SetAttributes(attribute.Int("http.status_code", http.StatusBadRequest))
 		w.WriteHeader(http.StatusBadRequest)
 		w.Write([]byte("URL should be in the pattern of /<subject>"))
 		return
 	}
 	subject := ss[1]
+	span.SetAttributes(attribute.String("nats.subject", subject))
 
 	fields := log.Fields{
 		"subject": subject,
 		"client":  r.RemoteAddr,
 	}
 	log.WithFields(fields).Info("Received HTTP request")
 
+	// Read request body with tracing
 	payload, err := io.ReadAll(r.Body)
 	if err != nil {
+		span.RecordError(err)
+		span.SetStatus(codes.Error, "Failed to read request body")
+		span.SetAttributes(attribute.Int("http.status_code", http.StatusInternalServerError))
 		w.WriteHeader(http.StatusInternalServerError)
 		_, err = fmt.Fprintf(w, "failed to read request body: %s", err)
 		if err != nil {
 			log.WithFields(fields).Errorf("failed to write payload: %v", err)
 		}
 		return
 	}
+	span.SetAttributes(attribute.Int("http.request.body_size", len(payload)))
 
 	// We can override the HTTP timeout by passing the `_timeout` query string
 	timeout := DEFAULT_HTTP_TIMEOUT
@@ -77,8 +107,9 @@ func (p *httpNatsProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 			timeout = DEFAULT_HTTP_TIMEOUT
 		}
 	}
+	span.SetAttributes(attribute.String("http.timeout", timeout.String()))
 
-	ctx, cancel := context.WithTimeout(r.Context(), timeout)
+	ctx, cancel := context.WithTimeout(ctx, timeout)
 	defer cancel()
 
 	url := strings.ReplaceAll(r.URL.String(), "/"+subject, "")
@@ -90,29 +121,61 @@ func (p *httpNatsProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		URL:     url,
 	})
 	if err != nil {
+		span.RecordError(err)
+		span.SetStatus(codes.Error, "Failed to encode message")
 		log.WithFields(fields).Errorf("failed to encode message: %v", err)
 		return
 	}
 
-	msg, err := p.nc.RequestWithContext(ctx, subject, body)
+	// Start a child span for the NATS request
+	ctx, natsSpan := tracer.Start(ctx, "nats.request",
+		trace.WithSpanKind(trace.SpanKindClient),
+		trace.WithAttributes(
+			attribute.String("nats.subject", subject),
+			attribute.Int("nats.message_size", len(body)),
+		),
+	)
+
+	// Inject trace context into NATS message headers
+	msg := nats.NewMsg(subject)
+	msg.Data = body
+	otel.GetTextMapPropagator().Inject(ctx, natsHeaderCarrier(msg.Header))
+
+	response, err := p.nc.RequestMsgWithContext(ctx, msg)
 	if err != nil {
+		natsSpan.RecordError(err)
+		natsSpan.SetStatus(codes.Error, "NATS request failed")
+		natsSpan.End()
+		span.SetStatus(codes.Error, "Service unavailable")
+		span.SetAttributes(attribute.Int("http.status_code", http.StatusServiceUnavailable))
 		w.WriteHeader(http.StatusServiceUnavailable)
 		w.Write([]byte(err.Error()))
 		return
 	}
+	natsSpan.SetAttributes(attribute.Int("nats.response_size", len(msg.Data)))
+	natsSpan.SetStatus(codes.Ok, "")
+	natsSpan.End()
 
 	rep := new(executor.Reply)
-	err = json.Unmarshal(msg.Data, rep)
+	err = json.Unmarshal(response.Data, rep)
 	if err != nil {
+		span.RecordError(err)
+		span.SetStatus(codes.Error, "Failed to unmarshal response")
+		span.SetAttributes(attribute.Int("http.status_code", http.StatusFailedDependency))
 		w.WriteHeader(http.StatusFailedDependency)
 		fmt.Fprintf(w, "Error: %v", err)
 		return
 	}
 
 	if rep.Error != "" {
+		span.SetAttributes(attribute.String("executor.error", rep.Error))
 		if rep.Error == (&executor.NoScriptFoundError{}).Error() {
+			span.SetStatus(codes.Error, "Script not found")
+			span.SetAttributes(attribute.Int("http.status_code", http.StatusNotFound))
 			w.WriteHeader(http.StatusNotFound)
 		} else {
+			span.SetStatus(codes.Error, rep.Error)
+			span.SetAttributes(attribute.Int("http.status_code", http.StatusInternalServerError))
 			w.WriteHeader(http.StatusInternalServerError)
 		}
 
@@ -126,6 +189,7 @@ func (p *httpNatsProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 
 	// Go through all the scripts to see if one is HTML
 	if t, sr := hasHTMLResult(rep.AllResults); t {
+		span.SetAttributes(attribute.Bool("response.is_html", true))
 		var hasContentType bool
 		for k, v := range sr.Headers {
 			if k == "Content-Type" {
@@ -136,27 +200,43 @@ func (p *httpNatsProxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 		if !hasContentType {
 			w.Header().Add("Content-Type", "text/html")
 		}
+		span.SetAttributes(
+			attribute.Int("http.status_code", sr.Code),
+			attribute.Int("http.response.body_size", len(sr.Payload)),
+		)
 		w.WriteHeader(sr.Code)
 
 		_, err = w.Write(sr.Payload)
 		if err != nil {
+			span.RecordError(err)
 			log.WithFields(fields).Errorf("failed to write reply back to HTTP response: %v", err)
 		}
 
+		span.SetStatus(codes.Ok, "")
 		// Since only the HTML page reply can "win" we ignore the rest
 		return
 	}
 
 	// Convert the results to bytes
+	span.SetAttributes(attribute.Bool("response.is_html", false))
 	rr, err := json.Marshal(rep.AllResults)
 	if err != nil {
+		span.RecordError(err)
 		log.WithFields(fields).Errorf("failed to serialize all results to JSON: %v", err)
 	}
 
+	span.SetAttributes(
+		attribute.Int("http.status_code", http.StatusOK),
+		attribute.Int("http.response.body_size", len(rr)),
+	)
+
 	_, err = w.Write(rr)
 	if err != nil {
+		span.RecordError(err)
 		log.WithFields(fields).Errorf("failed to write reply back to HTTP response: %v", err)
 	}
+
+	span.SetStatus(codes.Ok, "")
 }
 
 func hasHTMLResult(results map[string]*executor.ScriptResult) (bool, *executor.ScriptResult) {

cmd/server/main.go

@@ -2,25 +2,32 @@ package main
 
 import (
 	"encoding/json"
+	"errors"
 	"flag"
 	"os"
 	"os/signal"
 	"strings"
-	"syscall"
 	"time"
 
 	natsserver "github.com/nats-io/nats-server/v2/server"
 	"github.com/nats-io/nats.go"
 	log "github.com/sirupsen/logrus"
 	"golang.org/x/net/context"
 
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/codes"
+	"go.opentelemetry.io/otel/trace"
+
 	"github.com/numkem/msgscript"
 	"github.com/numkem/msgscript/executor"
 	msgplugin "github.com/numkem/msgscript/plugins"
 	msgstore "github.com/numkem/msgscript/store"
 )
 
 var version = "dev"
+var mainTracer = otel.Tracer("msgscript.main")
+
 func main() {
 	// Parse command-line flags
 	backendName := flag.String("backend", msgstore.BACKEND_FILE_NAME, "Storage backend to use (etcd, sqlite, flatfile)")
@@ -33,6 +40,9 @@ func main() {
 	scriptDir := flag.String("script", ".", "Script directory")
 	flag.Parse()
 
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt)
+	defer stop()
+
 	// Set up logging
 	level, err := log.ParseLevel(*logLevel)
 	if err != nil {
@@ -44,6 +54,20 @@ func main() {
 		log.SetLevel(log.DebugLevel)
 	}
 
+	if os.Getenv("TELEMETRY_TRACES") != "" {
+		log.WithField("kind", "traces").Info("Starting telemetry")
+
+		// Init traces
+		otelShutdown, err := setupOTelSDK(ctx)
+		if err != nil {
+			log.Errorf("failed to initialize opentelemetry traces: %v", err)
+			os.Exit(1)
+		}
+		defer func() {
+			err = errors.Join(err, otelShutdown(context.Background()))
+		}()
+	}
+
 	// Create the ScriptStore based on the selected backend
 	scriptStore, err := msgstore.StoreByName(*backendName, *etcdURL, *scriptDir, *libraryDir)
 	if err != nil {
@@ -104,15 +128,33 @@ func main() {
 
 	// Set up a message handler
 	_, err = nc.Subscribe(">", func(msg *nats.Msg) {
+		// Extract trace context from NATS message headers
+		ctx := otel.GetTextMapPropagator().Extract(
+			context.Background(),
+			natsHeaderCarrier(msg.Header),
+		)
+
+		// Start a span for the NATS message handling
+		ctx, span := mainTracer.Start(ctx, "nats.handle_message",
+			trace.WithSpanKind(trace.SpanKindServer),
+			trace.WithAttributes(
+				attribute.String("nats.subject", msg.Subject),
+				attribute.Int("nats.message_size", len(msg.Data)),
+			),
+		)
+		defer span.End()
+
 		log.Debugf("Received message on subject: %s", msg.Subject)
 
 		if strings.HasPrefix(msg.Subject, "_INBOX.") {
+			span.SetAttributes(attribute.Bool("nats.is_inbox", true))
+			span.SetStatus(codes.Ok, "Ignored inbox subject")
 			log.Debugf("Ignoring reply subject %s", msg.Subject)
 			return
 		}
 
 		m := new(executor.Message)
-		err = json.Unmarshal(msg.Data, m)
+		err := json.Unmarshal(msg.Data, m)
 		// if the payload isn't a JSON Message, take it as a whole
 		if err != nil {
 			m.Subject = msg.Subject
@@ -126,6 +168,12 @@ func main() {
 			"async":   m.Async,
 		}
 
+		span.SetAttributes(
+			attribute.Bool("message.raw", m.Raw),
+			attribute.Bool("message.async", m.Async),
+			attribute.String("message.executor", m.Executor),
+		)
+
 		// The above unmarshalling only applies to the structure of the JSON.
 		// Even if you feed it another JSON where none of the keys matches,
 		// it will just end up being an empty struct
@@ -139,37 +187,41 @@ func main() {
 		replier := &replier{nc: nc}
 		var messageReply executor.ReplyFunc
 		if m.Async {
+			span.SetAttributes(attribute.String("reply.mode", "async"))
 			messageReply = replier.AsyncReply(m, msg)
 			err = nc.Publish(msg.Reply, []byte("{}"))
 			if err != nil {
+				span.RecordError(err)
+				span.SetStatus(codes.Error, "Failed to publish async reply")
 				log.WithFields(fields).Errorf("failed to reply to message: %v", err)
 				return
 			}
 		} else {
+			span.SetAttributes(attribute.String("reply.mode", "sync"))
 			messageReply = replier.SyncReply(m, msg)
 		}
 
 		exec, err := executor.ExecutorByName(m.Executor, executors)
 		if err != nil {
+			span.RecordError(err)
+			span.SetStatus(codes.Error, "Failed to get executor")
 			log.WithError(err).Error("failed to get executor for message")
 			return
 		}
 
+		// Pass the context with trace info to the executor
 		exec.HandleMessage(ctx, m, messageReply)
+		span.SetStatus(codes.Ok, "Message handled")
 	})
 	if err != nil {
 		log.Fatalf("Failed to subscribe to NATS subjects: %v", err)
 	}
 
-	// Start HTTP Server
-	go runHTTP(*httpPort, *natsURL)
-
-	// Listen for system interrupts for graceful shutdown
-	sigChan := make(chan os.Signal, 1)
-	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)
-	<-sigChan
-	cancel()
+	defer func() {
+		log.Info("Received shutdown signal, stopping server...")
+		executor.StopAllExecutors(executors)
+	}()
 
-	log.Info("Received shutdown signal, stopping server...")
-	executor.StopAllExecutors(executors)
+	// Start HTTP Server
+	runHTTP(*httpPort, *natsURL)
 }