github.com/blend/go-sdk@v1.20220411.3/datadog/traceserver/server.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 traceserver
     9  
    10  import (
    11  	"context"
    12  	"fmt"
    13  	"log"
    14  	"net"
    15  	"net/http"
    16  	"strings"
    17  	"time"
    18  
    19  	msgp "github.com/tinylib/msgp/msgp"
    20  )
    21  
    22  // Server is a server for handling traces.
    23  type Server struct {
    24  	Addr     string
    25  	Log      *log.Logger
    26  	Listener net.Listener
    27  	Server   *http.Server
    28  	Handler  func(context.Context, ...*Span)
    29  }
    30  
    31  // Start starts the server.
    32  func (ts *Server) Start() error {
    33  	var err error
    34  	if ts.Handler == nil {
    35  		return fmt.Errorf("server cannot start; no handler provided")
    36  	}
    37  	if ts.Listener == nil && ts.Addr != "" {
    38  		ts.Listener, err = net.Listen("tcp", ts.Addr)
    39  		if err != nil {
    40  			return err
    41  		}
    42  	}
    43  	if ts.Listener == nil {
    44  		return fmt.Errorf("server cannot start; no listener or addr provided")
    45  	}
    46  
    47  	ts.logf("trace server listening: %s", ts.Listener.Addr().String())
    48  	ts.Server = &http.Server{
    49  		Handler: ts,
    50  	}
    51  	return ts.Server.Serve(ts.Listener)
    52  }
    53  
    54  // Stop stops the trace server.
    55  func (ts *Server) Stop() error {
    56  	if ts.Server == nil {
    57  		return nil
    58  	}
    59  	if err := ts.Server.Shutdown(context.Background()); err != nil {
    60  		return err
    61  	}
    62  	ts.Server = nil
    63  	return nil
    64  }
    65  
    66  func (ts *Server) ServeHTTP(rw http.ResponseWriter, req *http.Request) {
    67  	srw := &ResponseWriter{ResponseWriter: rw}
    68  
    69  	start := time.Now()
    70  	defer func() {
    71  		elapsed := time.Since(start)
    72  		ts.logf("%s %s %d %v %s", req.Method, req.URL.String(), srw.StatusCode, elapsed, FormatContentLength(srw.ContentLength))
    73  	}()
    74  
    75  	switch req.Method {
    76  	case http.MethodGet:
    77  		switch req.URL.Path {
    78  		case "/":
    79  			ts.handleGetIndex(srw, req)
    80  			return
    81  		default:
    82  		}
    83  	case http.MethodPost:
    84  		switch req.URL.Path {
    85  		case "/v0.4/traces":
    86  			ts.handlePostTraces(srw, req)
    87  			return
    88  		default:
    89  		}
    90  	default:
    91  	}
    92  	http.NotFound(srw, req)
    93  }
    94  
    95  //
    96  // handlers
    97  //
    98  
    99  func (ts *Server) handleGetIndex(rw http.ResponseWriter, req *http.Request) {
   100  	rw.WriteHeader(http.StatusOK)
   101  	fmt.Fprintf(rw, "Datadog Trace Echo")
   102  }
   103  
   104  func (ts *Server) handlePostTraces(rw http.ResponseWriter, req *http.Request) {
   105  	var payload SpanLists
   106  	if err := msgp.Decode(req.Body, &payload); err != nil {
   107  		http.Error(rw, err.Error(), http.StatusBadRequest)
   108  		return
   109  	}
   110  	for _, spanList := range payload {
   111  		ts.Handler(req.Context(), spanList...)
   112  	}
   113  	rw.WriteHeader(http.StatusOK)
   114  	fmt.Fprintf(rw, "OK!")
   115  }
   116  
   117  //
   118  // logging
   119  //
   120  
   121  func (ts *Server) logf(format string, args ...interface{}) {
   122  	if ts.Log != nil {
   123  		format = strings.TrimSpace(format)
   124  		ts.Log.Printf(format+"\n", args...)
   125  	}
   126  }
   127  
   128  func (ts *Server) logln(args ...interface{}) {
   129  	if ts.Log != nil {
   130  		ts.Log.Println(args...)
   131  	}
   132  }
   133  
   134  // FormatContentLength returns a string representation of a file size in bytes.
   135  func FormatContentLength(sizeBytes int) string {
   136  	if sizeBytes >= 1<<30 {
   137  		return fmt.Sprintf("%dgB", sizeBytes/(1<<30))
   138  	} else if sizeBytes >= 1<<20 {
   139  		return fmt.Sprintf("%dmB", sizeBytes/(1<<20))
   140  	} else if sizeBytes >= 1<<10 {
   141  		return fmt.Sprintf("%dkB", sizeBytes/(1<<10))
   142  	}
   143  	return fmt.Sprintf("%dB", sizeBytes)
   144  }