google.golang.org/grpc@v1.72.2/health/producer.go (about)

     1  /*
     2   *
     3   * Copyright 2024 gRPC authors.
     4   *
     5   * Licensed under the Apache License, Version 2.0 (the "License");
     6   * you may not use this file except in compliance with the License.
     7   * You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   *
    17   */
    18  
    19  package health
    20  
    21  import (
    22  	"context"
    23  	"sync"
    24  
    25  	"google.golang.org/grpc"
    26  	"google.golang.org/grpc/balancer"
    27  	"google.golang.org/grpc/codes"
    28  	"google.golang.org/grpc/connectivity"
    29  	"google.golang.org/grpc/internal"
    30  	"google.golang.org/grpc/status"
    31  )
    32  
    33  func init() {
    34  	producerBuilderSingleton = &producerBuilder{}
    35  	internal.RegisterClientHealthCheckListener = registerClientSideHealthCheckListener
    36  }
    37  
    38  type producerBuilder struct{}
    39  
    40  var producerBuilderSingleton *producerBuilder
    41  
    42  // Build constructs and returns a producer and its cleanup function.
    43  func (*producerBuilder) Build(cci any) (balancer.Producer, func()) {
    44  	p := &healthServiceProducer{
    45  		cc:     cci.(grpc.ClientConnInterface),
    46  		cancel: func() {},
    47  	}
    48  	return p, func() {
    49  		p.mu.Lock()
    50  		defer p.mu.Unlock()
    51  		p.cancel()
    52  	}
    53  }
    54  
    55  type healthServiceProducer struct {
    56  	// The following fields are initialized at build time and read-only after
    57  	// that and therefore do not need to be guarded by a mutex.
    58  	cc grpc.ClientConnInterface
    59  
    60  	mu     sync.Mutex
    61  	cancel func()
    62  }
    63  
    64  // registerClientSideHealthCheckListener accepts a listener to provide server
    65  // health state via the health service.
    66  func registerClientSideHealthCheckListener(ctx context.Context, sc balancer.SubConn, serviceName string, listener func(balancer.SubConnState)) func() {
    67  	pr, closeFn := sc.GetOrBuildProducer(producerBuilderSingleton)
    68  	p := pr.(*healthServiceProducer)
    69  	p.mu.Lock()
    70  	defer p.mu.Unlock()
    71  	p.cancel()
    72  	if listener == nil {
    73  		return closeFn
    74  	}
    75  
    76  	ctx, cancel := context.WithCancel(ctx)
    77  	p.cancel = cancel
    78  
    79  	go p.startHealthCheck(ctx, sc, serviceName, listener)
    80  	return closeFn
    81  }
    82  
    83  func (p *healthServiceProducer) startHealthCheck(ctx context.Context, sc balancer.SubConn, serviceName string, listener func(balancer.SubConnState)) {
    84  	newStream := func(method string) (any, error) {
    85  		return p.cc.NewStream(ctx, &grpc.StreamDesc{ServerStreams: true}, method)
    86  	}
    87  
    88  	setConnectivityState := func(state connectivity.State, err error) {
    89  		listener(balancer.SubConnState{
    90  			ConnectivityState: state,
    91  			ConnectionError:   err,
    92  		})
    93  	}
    94  
    95  	// Call the function through the internal variable as tests use it for
    96  	// mocking.
    97  	err := internal.HealthCheckFunc(ctx, newStream, setConnectivityState, serviceName)
    98  	if err == nil {
    99  		return
   100  	}
   101  	if status.Code(err) == codes.Unimplemented {
   102  		logger.Errorf("Subchannel health check is unimplemented at server side, thus health check is disabled for SubConn %p", sc)
   103  	} else {
   104  		logger.Errorf("Health checking failed for SubConn %p: %v", sc, err)
   105  	}
   106  }