github.com/apache/beam/sdks/v2@v2.48.2/go/container/tools/logging.go (about)

     1  // Licensed to the Apache Software Foundation (ASF) under one or more
     2  // contributor license agreements.  See the NOTICE file distributed with
     3  // this work for additional information regarding copyright ownership.
     4  // The ASF licenses this file to You under the Apache License, Version 2.0
     5  // (the "License"); you may not use this file except in compliance with
     6  // the License.  You may obtain a copy of the License at
     7  //
     8  //    http://www.apache.org/licenses/LICENSE-2.0
     9  //
    10  // Unless required by applicable law or agreed to in writing, software
    11  // distributed under the License is distributed on an "AS IS" BASIS,
    12  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  // See the License for the specific language governing permissions and
    14  // limitations under the License.
    15  
    16  package tools
    17  
    18  import (
    19  	"context"
    20  	"errors"
    21  	"fmt"
    22  	"log"
    23  	"os"
    24  	"sync"
    25  	"time"
    26  
    27  	fnpb "github.com/apache/beam/sdks/v2/go/pkg/beam/model/fnexecution_v1"
    28  	"github.com/apache/beam/sdks/v2/go/pkg/beam/util/grpcx"
    29  	"google.golang.org/protobuf/types/known/timestamppb"
    30  )
    31  
    32  // Logger is a wrapper around the FnAPI Logging Client, intended for
    33  // container boot loader use. Not intended for Beam end users.
    34  type Logger struct {
    35  	Endpoint string
    36  
    37  	client  logSender
    38  	closeFn func()
    39  	mu      sync.Mutex // To protect Send in the rare case multiple goroutines are calling this logger.
    40  }
    41  
    42  type logSender interface {
    43  	Send(*fnpb.LogEntry_List) error
    44  	CloseSend() error
    45  }
    46  
    47  // Close closes the grpc logging client.
    48  func (l *Logger) Close() {
    49  	if l.closeFn != nil {
    50  		l.client.CloseSend()
    51  		l.closeFn()
    52  		l.closeFn = nil
    53  		l.client = nil
    54  	}
    55  }
    56  
    57  // Log a message with the given severity.
    58  func (l *Logger) Log(ctx context.Context, sev fnpb.LogEntry_Severity_Enum, message string) {
    59  	l.mu.Lock()
    60  	defer l.mu.Unlock()
    61  
    62  	var exitErr error
    63  	defer func() {
    64  		if exitErr != nil {
    65  			log.Println("boot.go: error logging message over FnAPI. endpoint", l.Endpoint, "error:", exitErr, "message follows")
    66  			log.Println(sev.String(), message)
    67  		}
    68  	}()
    69  	if l.client == nil {
    70  		if l.Endpoint == "" {
    71  			exitErr = errors.New("no logging endpoint set")
    72  			return
    73  		}
    74  		cc, err := grpcx.Dial(ctx, l.Endpoint, 2*time.Minute)
    75  		if err != nil {
    76  			exitErr = err
    77  			return
    78  		}
    79  		l.closeFn = func() { cc.Close() }
    80  
    81  		l.client, err = fnpb.NewBeamFnLoggingClient(cc).Logging(ctx)
    82  		if err != nil {
    83  			exitErr = err
    84  			l.Close()
    85  			return
    86  		}
    87  	}
    88  
    89  	err := l.client.Send(&fnpb.LogEntry_List{
    90  		LogEntries: []*fnpb.LogEntry{
    91  			{
    92  				Severity:  sev,
    93  				Timestamp: timestamppb.Now(),
    94  				Message:   message,
    95  			},
    96  		},
    97  	})
    98  	if err != nil {
    99  		exitErr = err
   100  		return
   101  	}
   102  }
   103  
   104  // Printf logs the message with Debug severity.
   105  func (l *Logger) Printf(ctx context.Context, format string, args ...any) {
   106  	l.Log(ctx, fnpb.LogEntry_Severity_DEBUG, fmt.Sprintf(format, args...))
   107  }
   108  
   109  // Warnf logs the message with Warning severity.
   110  func (l *Logger) Warnf(ctx context.Context, format string, args ...any) {
   111  	l.Log(ctx, fnpb.LogEntry_Severity_WARN, fmt.Sprintf(format, args...))
   112  }
   113  
   114  // Fatalf logs the message with Critical severity, and then calls os.Exit(1).
   115  func (l *Logger) Fatalf(ctx context.Context, format string, args ...any) {
   116  	l.Log(ctx, fnpb.LogEntry_Severity_CRITICAL, fmt.Sprintf(format, args...))
   117  	os.Exit(1)
   118  }