| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202 |
- // Package grpcclient provides a Go client for the broad-announce ingestd
- // gRPC service (SPEC §19). It wraps the generated pb.IngestClient and
- // adds Bearer-token auth, stream-level retry semantics, and ergonomic
- // Send/Recv primitives for use by any peer service that publishes alerts.
- package grpcclient
- import (
- "context"
- "fmt"
- "io"
- "time"
- pbv1 "git3.techno-world.net/lrosales/broad-announce/gen/go/broadannounce/v1"
- "google.golang.org/grpc"
- "google.golang.org/grpc/credentials"
- "google.golang.org/grpc/credentials/insecure"
- "google.golang.org/grpc/keepalive"
- "google.golang.org/grpc/metadata"
- "google.golang.org/grpc/status"
- )
- // Client is a gRPC client for the Ingest service. It is safe for
- // concurrent use by multiple goroutines.
- type Client struct {
- conn *grpc.ClientConn
- client pbv1.IngestClient
- apiKey string
- maxRetries int
- keepalive keepalive.ClientParameters
- insecure bool
- }
- // Stream represents an open StreamAlerts bidirectional stream. A single
- // Stream is not safe for concurrent use — use one stream per producer
- // goroutine, or add external synchronization.
- type Stream struct {
- stream pbv1.Ingest_StreamAlertsClient
- }
- // New connects to the Ingest service at addr and returns a Client.
- // It does not perform any I/O; the connection is established lazily on
- // the first call to Stream.
- //
- // - addr: host:port of the ingestd gRPC server (default :9090)
- // - opts: functional options (WithAPIKey, WithMaxRetries,
- // WithKeepalive, WithInsecure)
- //
- // Production deployments should use TLS or mTLS. The default TLS
- // credentials use the host's root CA set. For mTLS, pass a custom
- // grpc.DialOption built with credentials.NewTLS(tlsConfig) via a
- // private Option (or replace WithInsecure in development only).
- func New(addr string, opts ...Option) (*Client, error) {
- c := &Client{maxRetries: 3}
- for _, opt := range opts {
- opt(c)
- }
- var dialOpts []grpc.DialOption
- if c.insecure {
- dialOpts = append(dialOpts, grpc.WithTransportCredentials(insecure.NewCredentials()))
- } else {
- // Secure TLS. Replace by passing a custom tls.Config with
- // ClientCAs (for mTLS) or ServerName override (for TLS).
- dialOpts = append(dialOpts,
- grpc.WithTransportCredentials(credentials.NewTLS(nil)))
- }
- if c.keepalive.Time > 0 || c.keepalive.Timeout > 0 {
- dialOpts = append(dialOpts, grpc.WithKeepaliveParams(c.keepalive))
- }
- // Unary interceptor: attach API key to every RPC.
- dialOpts = append(dialOpts, grpc.WithUnaryInterceptor(
- func(ctx context.Context, method string, req, reply interface{},
- cc *grpc.ClientConn, invoker grpc.UnaryInvoker, _ ...grpc.CallOption) error {
- if c.apiKey != "" {
- ctx = metadata.AppendToOutgoingContext(ctx,
- "authorization", "Bearer "+c.apiKey)
- }
- return invoker(ctx, method, req, reply, cc)
- }))
- conn, err := grpc.NewClient(addr, dialOpts...)
- if err != nil {
- return nil, fmt.Errorf("grpc.Dial(%q): %w", addr, err)
- }
- c.conn = conn
- c.client = pbv1.NewIngestClient(conn)
- return c, nil
- }
- // Stream starts a StreamAlerts bidirectional stream. The caller must call
- // CloseSend when done sending. The context is used for stream deadlines.
- func (c *Client) Stream(ctx context.Context) (*Stream, error) {
- if c.apiKey != "" {
- md := metadata.Pairs("authorization", "Bearer "+c.apiKey)
- ctx = metadata.NewOutgoingContext(ctx, md)
- }
- stream, err := c.client.StreamAlerts(ctx)
- if err != nil {
- return nil, fmt.Errorf("StreamAlerts: %w", err)
- }
- return &Stream{stream: stream}, nil
- }
- // Send delivers one alert to the server. It blocks until the server
- // acknowledges receipt (not processing — that is async). Send retries on
- // RATE_LIMITED / UNAVAILABLE up to c.maxRetries, honouring the
- // x-retry-after-ms response header when set.
- func (s *Stream) Send(ctx context.Context, alert *pbv1.Alert) error {
- var lastErr error
- for attempt := 0; attempt <= 3; attempt++ {
- select {
- case <-ctx.Done():
- return ctx.Err()
- default:
- }
- lastErr = s.stream.Send(alert)
- if lastErr == nil {
- return nil
- }
- st, ok := status.FromError(lastErr)
- if !ok || !retryable(st.Code()) {
- return lastErr
- }
- if attempt < 3 {
- delay := retryDelay(s.stream.Context(), attempt)
- select {
- case <-ctx.Done():
- return ctx.Err()
- case <-time.After(delay):
- }
- }
- }
- return fmt.Errorf("Stream.Send (retries exhausted): %w", lastErr)
- }
- // Recv returns the next Ack from the server. It blocks until an Ack is
- // available or the stream is closed. It retries on UNAVAILABLE up to
- // c.maxRetries. Returns io.EOF when the server closes the stream.
- func (s *Stream) Recv(ctx context.Context) (*pbv1.Ack, error) {
- var lastErr error
- for attempt := 0; attempt <= 3; attempt++ {
- select {
- case <-ctx.Done():
- return nil, ctx.Err()
- default:
- }
- ack, err := s.stream.Recv()
- if err == nil {
- return ack, nil
- }
- lastErr = err
- if err == io.EOF {
- return nil, io.EOF
- }
- st, ok := status.FromError(err)
- if !ok || !retryable(st.Code()) {
- return nil, err
- }
- if attempt < 3 {
- delay := retryDelay(s.stream.Context(), attempt)
- select {
- case <-ctx.Done():
- return nil, ctx.Err()
- case <-time.After(delay):
- }
- }
- }
- return nil, fmt.Errorf("Stream.Recv (retries exhausted): %w", lastErr)
- }
- // CloseSend signals that the caller is done sending alerts. The server
- // will half-close its side of the stream; Recv continues to work until
- // the server closes the stream or the context is cancelled.
- func (s *Stream) CloseSend() error { return s.stream.CloseSend() }
- // Close tears down the underlying gRPC connection. Call it when the
- // Client is no longer needed.
- func (c *Client) Close() error { return c.conn.Close() }
- // retryDelay returns the sleep interval for a given retry attempt,
- // checking the x-retry-after-ms header for a server-specified value.
- func retryDelay(streamCtx context.Context, attempt int) time.Duration {
- if md, ok := metadata.FromIncomingContext(streamCtx); ok {
- if vals := md.Get("x-retry-after-ms"); len(vals) > 0 {
- var ms int
- fmt.Sscanf(vals[0], "%d", &ms)
- if ms > 0 {
- return time.Duration(ms) * time.Millisecond
- }
- }
- }
- return time.Duration(100*attempt) * time.Millisecond
- }
|