mirror of
https://github.com/pomerium/pomerium.git
synced 2025-04-28 18:06:34 +02:00
103 lines
2.7 KiB
Go
103 lines
2.7 KiB
Go
package trace
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"fmt"
|
|
"net"
|
|
"time"
|
|
|
|
"github.com/pomerium/pomerium/internal/log"
|
|
coltracepb "go.opentelemetry.io/proto/otlp/collector/trace/v1"
|
|
"google.golang.org/grpc"
|
|
"google.golang.org/grpc/credentials/insecure"
|
|
"google.golang.org/grpc/test/bufconn"
|
|
|
|
"go.opentelemetry.io/otel/exporters/otlp/otlptrace"
|
|
"go.opentelemetry.io/otel/exporters/otlp/otlptrace/otlptracegrpc"
|
|
)
|
|
|
|
const localExporterMetadataKey = "x-local-exporter"
|
|
|
|
// Export implements ptraceotlp.GRPCServer.
|
|
func (srv *ExporterServer) Export(ctx context.Context, req *coltracepb.ExportTraceServiceRequest) (*coltracepb.ExportTraceServiceResponse, error) {
|
|
if srv.observer != nil {
|
|
srv.observeExport(ctx, req)
|
|
}
|
|
if err := srv.remoteClient.UploadTraces(ctx, req.GetResourceSpans()); err != nil {
|
|
log.Ctx(ctx).Err(err).Msg("error uploading traces")
|
|
return nil, err
|
|
}
|
|
return &coltracepb.ExportTraceServiceResponse{}, nil
|
|
}
|
|
|
|
type ExporterServer struct {
|
|
coltracepb.UnimplementedTraceServiceServer
|
|
server *grpc.Server
|
|
observer *spanObserver
|
|
remoteClient otlptrace.Client
|
|
cc *grpc.ClientConn
|
|
}
|
|
|
|
func NewServer(ctx context.Context) *ExporterServer {
|
|
sys := systemContextFromContext(ctx)
|
|
ex := &ExporterServer{
|
|
remoteClient: sys.remoteClient,
|
|
observer: sys.observer,
|
|
server: grpc.NewServer(grpc.Creds(insecure.NewCredentials())),
|
|
}
|
|
coltracepb.RegisterTraceServiceServer(ex.server, ex)
|
|
return ex
|
|
}
|
|
|
|
func (srv *ExporterServer) Start(ctx context.Context) {
|
|
lis := bufconn.Listen(2 * 1024 * 1024)
|
|
go func() {
|
|
if err := srv.remoteClient.Start(ctx); err != nil {
|
|
if !errors.Is(err, ErrNoClient) {
|
|
panic(fmt.Errorf("bug: %w", err))
|
|
}
|
|
}
|
|
_ = srv.server.Serve(lis)
|
|
}()
|
|
cc, err := grpc.NewClient("passthrough://ignore",
|
|
grpc.WithContextDialer(func(context.Context, string) (net.Conn, error) {
|
|
return lis.Dial()
|
|
}), grpc.WithTransportCredentials(insecure.NewCredentials()))
|
|
if err != nil {
|
|
panic(err)
|
|
}
|
|
srv.cc = cc
|
|
}
|
|
|
|
func (srv *ExporterServer) NewClient() otlptrace.Client {
|
|
return otlptracegrpc.NewClient(
|
|
otlptracegrpc.WithGRPCConn(srv.cc),
|
|
otlptracegrpc.WithTimeout(1*time.Minute),
|
|
otlptracegrpc.WithHeaders(map[string]string{
|
|
localExporterMetadataKey: "1",
|
|
}),
|
|
)
|
|
}
|
|
|
|
func (srv *ExporterServer) Shutdown(ctx context.Context) error {
|
|
stopped := make(chan struct{})
|
|
go func() {
|
|
srv.server.GracefulStop()
|
|
close(stopped)
|
|
}()
|
|
select {
|
|
case <-stopped:
|
|
case <-ctx.Done():
|
|
return context.Cause(ctx)
|
|
}
|
|
var errs []error
|
|
if err := WaitForSpans(ctx, 30*time.Second); err != nil {
|
|
errs = append(errs, err)
|
|
}
|
|
if err := srv.remoteClient.Stop(ctx); err != nil {
|
|
errs = append(errs, err)
|
|
}
|
|
srv.cc.Close()
|
|
return errors.Join(errs...)
|
|
}
|