github.com/minio/minio@v0.0.0-20240328213742-3f72439b8a27/internal/grid/handlers.go (about)

     1  // Copyright (c) 2015-2023 MinIO, Inc.
     2  //
     3  // This file is part of MinIO Object Storage stack
     4  //
     5  // This program is free software: you can redistribute it and/or modify
     6  // it under the terms of the GNU Affero General Public License as published by
     7  // the Free Software Foundation, either version 3 of the License, or
     8  // (at your option) any later version.
     9  //
    10  // This program is distributed in the hope that it will be useful
    11  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    12  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    13  // GNU Affero General Public License for more details.
    14  //
    15  // You should have received a copy of the GNU Affero General Public License
    16  // along with this program.  If not, see <http://www.gnu.org/licenses/>.
    17  
    18  package grid
    19  
    20  import (
    21  	"context"
    22  	"encoding/hex"
    23  	"errors"
    24  	"fmt"
    25  	"strings"
    26  	"sync"
    27  
    28  	"github.com/minio/minio/internal/hash/sha256"
    29  	xioutil "github.com/minio/minio/internal/ioutil"
    30  	"github.com/minio/minio/internal/logger"
    31  	"github.com/tinylib/msgp/msgp"
    32  )
    33  
    34  //go:generate stringer -type=HandlerID -output=handlers_string.go -trimprefix=Handler msg.go $GOFILE
    35  
    36  // HandlerID is a handler identifier.
    37  // It is used to determine request routing on the server.
    38  // Handlers can be registered with a static subroute.
    39  // Do NOT remove or change the order of existing handlers.
    40  const (
    41  	// handlerInvalid is reserved to check for uninitialized values.
    42  	handlerInvalid HandlerID = iota
    43  	HandlerLockLock
    44  	HandlerLockRLock
    45  	HandlerLockUnlock
    46  	HandlerLockRUnlock
    47  	HandlerLockRefresh
    48  	HandlerLockForceUnlock
    49  	HandlerWalkDir
    50  	HandlerStatVol
    51  	HandlerDiskInfo
    52  	HandlerNSScanner
    53  	HandlerReadXL
    54  	HandlerReadVersion
    55  	HandlerDeleteFile
    56  	HandlerDeleteVersion
    57  	HandlerUpdateMetadata
    58  	HandlerWriteMetadata
    59  	HandlerCheckParts
    60  	HandlerRenameData
    61  	HandlerRenameFile
    62  	HandlerReadAll
    63  	HandlerServerVerify
    64  	HandlerTrace
    65  	HandlerListen
    66  	HandlerDeleteBucketMetadata
    67  	HandlerLoadBucketMetadata
    68  	HandlerReloadSiteReplicationConfig
    69  	HandlerReloadPoolMeta
    70  	HandlerStopRebalance
    71  	HandlerLoadRebalanceMeta
    72  	HandlerLoadTransitionTierConfig
    73  	HandlerDeletePolicy
    74  	HandlerLoadPolicy
    75  	HandlerLoadPolicyMapping
    76  	HandlerDeleteServiceAccount
    77  	HandlerLoadServiceAccount
    78  	HandlerDeleteUser
    79  	HandlerLoadUser
    80  	HandlerLoadGroup
    81  	HandlerHealBucket
    82  	HandlerMakeBucket
    83  	HandlerHeadBucket
    84  	HandlerDeleteBucket
    85  	HandlerGetMetrics
    86  	HandlerGetResourceMetrics
    87  	HandlerGetMemInfo
    88  	HandlerGetProcInfo
    89  	HandlerGetOSInfo
    90  	HandlerGetPartitions
    91  	HandlerGetNetInfo
    92  	HandlerGetCPUs
    93  	HandlerServerInfo
    94  	HandlerGetSysConfig
    95  	HandlerGetSysServices
    96  	HandlerGetSysErrors
    97  	HandlerGetAllBucketStats
    98  	HandlerGetBucketStats
    99  	HandlerGetSRMetrics
   100  	HandlerGetPeerMetrics
   101  	HandlerGetMetacacheListing
   102  	HandlerUpdateMetacacheListing
   103  	HandlerGetPeerBucketMetrics
   104  	HandlerStorageInfo
   105  	HandlerConsoleLog
   106  	HandlerListDir
   107  	HandlerGetLocks
   108  	HandlerBackgroundHealStatus
   109  	HandlerGetLastDayTierStats
   110  	HandlerSignalService
   111  	HandlerGetBandwidth
   112  	HandlerWriteAll
   113  	HandlerListBuckets
   114  
   115  	// Add more above here ^^^
   116  	// If all handlers are used, the type of Handler can be changed.
   117  	// Handlers have no versioning, so non-compatible handler changes must result in new IDs.
   118  	handlerTest
   119  	handlerTest2
   120  	handlerLast
   121  )
   122  
   123  // handlerPrefixes are prefixes for handler IDs used for tracing.
   124  // If a handler is not listed here, it will be traced with "grid" prefix.
   125  var handlerPrefixes = [handlerLast]string{
   126  	HandlerLockLock:                    lockPrefix,
   127  	HandlerLockRLock:                   lockPrefix,
   128  	HandlerLockUnlock:                  lockPrefix,
   129  	HandlerLockRUnlock:                 lockPrefix,
   130  	HandlerLockRefresh:                 lockPrefix,
   131  	HandlerLockForceUnlock:             lockPrefix,
   132  	HandlerWalkDir:                     storagePrefix,
   133  	HandlerStatVol:                     storagePrefix,
   134  	HandlerDiskInfo:                    storagePrefix,
   135  	HandlerNSScanner:                   storagePrefix,
   136  	HandlerReadXL:                      storagePrefix,
   137  	HandlerReadVersion:                 storagePrefix,
   138  	HandlerDeleteFile:                  storagePrefix,
   139  	HandlerDeleteVersion:               storagePrefix,
   140  	HandlerUpdateMetadata:              storagePrefix,
   141  	HandlerWriteMetadata:               storagePrefix,
   142  	HandlerCheckParts:                  storagePrefix,
   143  	HandlerRenameData:                  storagePrefix,
   144  	HandlerRenameFile:                  storagePrefix,
   145  	HandlerReadAll:                     storagePrefix,
   146  	HandlerWriteAll:                    storagePrefix,
   147  	HandlerServerVerify:                bootstrapPrefix,
   148  	HandlerTrace:                       peerPrefix,
   149  	HandlerListen:                      peerPrefix,
   150  	HandlerDeleteBucketMetadata:        peerPrefix,
   151  	HandlerLoadBucketMetadata:          peerPrefix,
   152  	HandlerReloadSiteReplicationConfig: peerPrefix,
   153  	HandlerReloadPoolMeta:              peerPrefix,
   154  	HandlerStopRebalance:               peerPrefix,
   155  	HandlerLoadRebalanceMeta:           peerPrefix,
   156  	HandlerLoadTransitionTierConfig:    peerPrefix,
   157  	HandlerDeletePolicy:                peerPrefix,
   158  	HandlerLoadPolicy:                  peerPrefix,
   159  	HandlerLoadPolicyMapping:           peerPrefix,
   160  	HandlerDeleteServiceAccount:        peerPrefix,
   161  	HandlerLoadServiceAccount:          peerPrefix,
   162  	HandlerDeleteUser:                  peerPrefix,
   163  	HandlerLoadUser:                    peerPrefix,
   164  	HandlerLoadGroup:                   peerPrefix,
   165  	HandlerMakeBucket:                  peerPrefixS3,
   166  	HandlerHeadBucket:                  peerPrefixS3,
   167  	HandlerDeleteBucket:                peerPrefixS3,
   168  	HandlerHealBucket:                  healPrefix,
   169  	HandlerGetMetrics:                  peerPrefix,
   170  	HandlerGetResourceMetrics:          peerPrefix,
   171  	HandlerGetMemInfo:                  peerPrefix,
   172  	HandlerGetProcInfo:                 peerPrefix,
   173  	HandlerGetOSInfo:                   peerPrefix,
   174  	HandlerGetPartitions:               peerPrefix,
   175  	HandlerGetNetInfo:                  peerPrefix,
   176  	HandlerGetCPUs:                     peerPrefix,
   177  	HandlerServerInfo:                  peerPrefix,
   178  	HandlerGetSysConfig:                peerPrefix,
   179  	HandlerGetSysServices:              peerPrefix,
   180  	HandlerGetSysErrors:                peerPrefix,
   181  	HandlerGetAllBucketStats:           peerPrefix,
   182  	HandlerGetBucketStats:              peerPrefix,
   183  	HandlerGetSRMetrics:                peerPrefix,
   184  	HandlerGetPeerMetrics:              peerPrefix,
   185  	HandlerGetMetacacheListing:         peerPrefix,
   186  	HandlerUpdateMetacacheListing:      peerPrefix,
   187  	HandlerGetPeerBucketMetrics:        peerPrefix,
   188  	HandlerStorageInfo:                 peerPrefix,
   189  	HandlerConsoleLog:                  peerPrefix,
   190  	HandlerListDir:                     storagePrefix,
   191  	HandlerListBuckets:                 peerPrefixS3,
   192  }
   193  
   194  const (
   195  	lockPrefix      = "lockR"
   196  	storagePrefix   = "storageR"
   197  	bootstrapPrefix = "bootstrap"
   198  	peerPrefix      = "peer"
   199  	peerPrefixS3    = "peerS3"
   200  	healPrefix      = "heal"
   201  )
   202  
   203  func init() {
   204  	// Static check if we exceed 255 handler ids.
   205  	// Extend the type to uint16 when hit.
   206  	if handlerLast > 255 {
   207  		panic(fmt.Sprintf("out of handler IDs. %d > %d", handlerLast, 255))
   208  	}
   209  }
   210  
   211  func (h HandlerID) valid() bool {
   212  	return h != handlerInvalid && h < handlerLast
   213  }
   214  
   215  func (h HandlerID) isTestHandler() bool {
   216  	return h >= handlerTest && h <= handlerTest2
   217  }
   218  
   219  // RemoteErr is a remote error type.
   220  // Any error seen on a remote will be returned like this.
   221  type RemoteErr string
   222  
   223  // NewRemoteErr creates a new remote error.
   224  // The error type is not preserved.
   225  func NewRemoteErr(err error) *RemoteErr {
   226  	if err == nil {
   227  		return nil
   228  	}
   229  	r := RemoteErr(err.Error())
   230  	return &r
   231  }
   232  
   233  // NewRemoteErrf creates a new remote error from a format string.
   234  func NewRemoteErrf(format string, a ...any) *RemoteErr {
   235  	r := RemoteErr(fmt.Sprintf(format, a...))
   236  	return &r
   237  }
   238  
   239  // NewNPErr is a helper to no payload and optional remote error.
   240  // The error type is not preserved.
   241  func NewNPErr(err error) (NoPayload, *RemoteErr) {
   242  	if err == nil {
   243  		return NoPayload{}, nil
   244  	}
   245  	r := RemoteErr(err.Error())
   246  	return NoPayload{}, &r
   247  }
   248  
   249  // NewRemoteErrString creates a new remote error from a string.
   250  func NewRemoteErrString(msg string) *RemoteErr {
   251  	r := RemoteErr(msg)
   252  	return &r
   253  }
   254  
   255  func (r RemoteErr) Error() string {
   256  	return string(r)
   257  }
   258  
   259  // Is returns if the string representation matches.
   260  func (r *RemoteErr) Is(other error) bool {
   261  	if r == nil || other == nil {
   262  		return r == other
   263  	}
   264  	var o RemoteErr
   265  	if errors.As(other, &o) {
   266  		return r == &o
   267  	}
   268  	return false
   269  }
   270  
   271  // IsRemoteErr returns the value if the error is a RemoteErr.
   272  func IsRemoteErr(err error) *RemoteErr {
   273  	var r RemoteErr
   274  	if errors.As(err, &r) {
   275  		return &r
   276  	}
   277  	return nil
   278  }
   279  
   280  type (
   281  	// SingleHandlerFn is handlers for one to one requests.
   282  	// A non-nil error value will be returned as RemoteErr(msg) to client.
   283  	// No client information or cancellation (deadline) is available.
   284  	// Include this in payload if needed.
   285  	// Payload should be recycled with PutByteBuffer if not needed after the call.
   286  	SingleHandlerFn func(payload []byte) ([]byte, *RemoteErr)
   287  
   288  	// StatelessHandlerFn must handle incoming stateless request.
   289  	// A non-nil error value will be returned as RemoteErr(msg) to client.
   290  	StatelessHandlerFn func(ctx context.Context, payload []byte, resp chan<- []byte) *RemoteErr
   291  
   292  	// StatelessHandler is handlers for one to many requests,
   293  	// where responses may be dropped.
   294  	// Stateless requests provide no incoming stream and there is no flow control
   295  	// on outgoing messages.
   296  	StatelessHandler struct {
   297  		Handle StatelessHandlerFn
   298  		// OutCapacity is the output capacity on the caller.
   299  		// If <= 0 capacity will be 1.
   300  		OutCapacity int
   301  	}
   302  
   303  	// StreamHandlerFn must process a request with an optional initial payload.
   304  	// It must keep consuming from 'in' until it returns.
   305  	// 'in' and 'out' are independent.
   306  	// The handler should never close out.
   307  	// Buffers received from 'in'  can be recycled with PutByteBuffer.
   308  	// Buffers sent on out can not be referenced once sent.
   309  	StreamHandlerFn func(ctx context.Context, payload []byte, in <-chan []byte, out chan<- []byte) *RemoteErr
   310  
   311  	// StreamHandler handles fully bidirectional streams,
   312  	// There is flow control in both directions.
   313  	StreamHandler struct {
   314  		// Handle an incoming request. Initial payload is sent.
   315  		// Additional input packets (if any) are streamed to request.
   316  		// Upstream will block when request channel is full.
   317  		// Response packets can be sent at any time.
   318  		// Any non-nil error sent as response means no more responses are sent.
   319  		Handle StreamHandlerFn
   320  
   321  		// Subroute for handler.
   322  		// Subroute must be static and clients should specify a matching subroute.
   323  		// Should not be set unless there are different handlers for the same HandlerID.
   324  		Subroute string
   325  
   326  		// OutCapacity is the output capacity. If <= 0 capacity will be 1.
   327  		OutCapacity int
   328  
   329  		// InCapacity is the output capacity.
   330  		// If == 0 no input is expected
   331  		InCapacity int
   332  	}
   333  )
   334  
   335  type subHandlerID [32]byte
   336  
   337  func makeSubHandlerID(id HandlerID, subRoute string) subHandlerID {
   338  	b := subHandlerID(sha256.Sum256([]byte(subRoute)))
   339  	b[0] = byte(id)
   340  	b[1] = 0 // Reserved
   341  	return b
   342  }
   343  
   344  func (s subHandlerID) withHandler(id HandlerID) subHandlerID {
   345  	s[0] = byte(id)
   346  	s[1] = 0 // Reserved
   347  	return s
   348  }
   349  
   350  func (s *subHandlerID) String() string {
   351  	if s == nil {
   352  		return ""
   353  	}
   354  	return hex.EncodeToString(s[:])
   355  }
   356  
   357  func makeZeroSubHandlerID(id HandlerID) subHandlerID {
   358  	return subHandlerID{byte(id)}
   359  }
   360  
   361  type handlers struct {
   362  	single    [handlerLast]SingleHandlerFn
   363  	stateless [handlerLast]*StatelessHandler
   364  	streams   [handlerLast]*StreamHandler
   365  
   366  	subSingle    map[subHandlerID]SingleHandlerFn
   367  	subStateless map[subHandlerID]*StatelessHandler
   368  	subStreams   map[subHandlerID]*StreamHandler
   369  }
   370  
   371  func (h *handlers) init() {
   372  	h.subSingle = make(map[subHandlerID]SingleHandlerFn)
   373  	h.subStateless = make(map[subHandlerID]*StatelessHandler)
   374  	h.subStreams = make(map[subHandlerID]*StreamHandler)
   375  }
   376  
   377  func (h *handlers) hasAny(id HandlerID) bool {
   378  	if !id.valid() {
   379  		return false
   380  	}
   381  	return h.single[id] != nil || h.stateless[id] != nil || h.streams[id] != nil
   382  }
   383  
   384  func (h *handlers) hasSubhandler(id subHandlerID) bool {
   385  	return h.subSingle[id] != nil || h.subStateless[id] != nil || h.subStreams[id] != nil
   386  }
   387  
   388  // RoundTripper provides an interface for type roundtrip serialization.
   389  type RoundTripper interface {
   390  	msgp.Unmarshaler
   391  	msgp.Marshaler
   392  	msgp.Sizer
   393  
   394  	comparable
   395  }
   396  
   397  // SingleHandler is a type safe handler for single roundtrip requests.
   398  type SingleHandler[Req, Resp RoundTripper] struct {
   399  	id            HandlerID
   400  	sharedResp    bool
   401  	callReuseReq  bool
   402  	ignoreNilConn bool
   403  
   404  	newReq  func() Req
   405  	newResp func() Resp
   406  
   407  	recycleReq  func(Req)
   408  	recycleResp func(Resp)
   409  }
   410  
   411  func recycleFunc[RT RoundTripper](newRT func() RT) (newFn func() RT, recycle func(r RT)) {
   412  	rAny := any(newRT())
   413  	var rZero RT
   414  	if _, ok := rAny.(Recycler); ok {
   415  		return newRT, func(r RT) {
   416  			if r != rZero {
   417  				if rc, ok := any(r).(Recycler); ok {
   418  					rc.Recycle()
   419  				}
   420  			}
   421  		}
   422  	}
   423  	pool := sync.Pool{
   424  		New: func() interface{} {
   425  			return newRT()
   426  		},
   427  	}
   428  	return func() RT { return pool.Get().(RT) },
   429  		func(r RT) {
   430  			if r != rZero {
   431  				pool.Put(r)
   432  			}
   433  		}
   434  }
   435  
   436  // NewSingleHandler creates a typed handler that can provide Marshal/Unmarshal.
   437  // Use Register to register a server handler.
   438  // Use Call to initiate a clientside call.
   439  func NewSingleHandler[Req, Resp RoundTripper](h HandlerID, newReq func() Req, newResp func() Resp) *SingleHandler[Req, Resp] {
   440  	s := SingleHandler[Req, Resp]{id: h}
   441  	s.newReq, s.recycleReq = recycleFunc[Req](newReq)
   442  	s.newResp, s.recycleResp = recycleFunc[Resp](newResp)
   443  	if _, ok := any(newReq()).(Recycler); ok {
   444  		s.callReuseReq = true
   445  	}
   446  	return &s
   447  }
   448  
   449  // PutResponse will accept a response for reuse.
   450  // This can be used by a caller to recycle a response after receiving it from a Call.
   451  func (h *SingleHandler[Req, Resp]) PutResponse(r Resp) {
   452  	h.recycleResp(r)
   453  }
   454  
   455  // AllowCallRequestPool indicates it is safe to reuse the request
   456  // on the client side, meaning the request is recycled/pooled when a request is sent.
   457  // CAREFUL: This should only be used when there are no pointers, slices that aren't freshly constructed.
   458  func (h *SingleHandler[Req, Resp]) AllowCallRequestPool(b bool) *SingleHandler[Req, Resp] {
   459  	h.callReuseReq = b
   460  	return h
   461  }
   462  
   463  // IgnoreNilConn will ignore nil connections when calling.
   464  // This will make Call return nil instead of ErrDisconnected when the connection is nil.
   465  // This may only be set ONCE before use.
   466  func (h *SingleHandler[Req, Resp]) IgnoreNilConn() *SingleHandler[Req, Resp] {
   467  	if h.ignoreNilConn {
   468  		logger.LogOnceIf(context.Background(), fmt.Errorf("%s: IgnoreNilConn called twice", h.id.String()), h.id.String()+"IgnoreNilConn")
   469  	}
   470  	h.ignoreNilConn = true
   471  	return h
   472  }
   473  
   474  // WithSharedResponse indicates it is unsafe to reuse the response
   475  // when it has been returned on a handler.
   476  // This will disable automatic response recycling/pooling.
   477  // Typically this is used when the response sharing part of its data structure.
   478  func (h *SingleHandler[Req, Resp]) WithSharedResponse() *SingleHandler[Req, Resp] {
   479  	h.sharedResp = true
   480  	return h
   481  }
   482  
   483  // NewResponse creates a new response.
   484  // Handlers can use this to create a reusable response.
   485  // The response may be reused, so caller should clear any fields.
   486  func (h *SingleHandler[Req, Resp]) NewResponse() Resp {
   487  	return h.newResp()
   488  }
   489  
   490  // NewRequest creates a new request.
   491  // Handlers can use this to create a reusable request.
   492  // The request may be reused, so caller should clear any fields.
   493  func (h *SingleHandler[Req, Resp]) NewRequest() Req {
   494  	return h.newReq()
   495  }
   496  
   497  // Register a handler for a Req -> Resp roundtrip.
   498  // Requests are automatically recycled.
   499  func (h *SingleHandler[Req, Resp]) Register(m *Manager, handle func(req Req) (resp Resp, err *RemoteErr), subroute ...string) error {
   500  	if h.newReq == nil {
   501  		return errors.New("newReq nil in NewSingleHandler")
   502  	}
   503  	if h.newResp == nil {
   504  		return errors.New("newResp nil in NewSingleHandler")
   505  	}
   506  	return m.RegisterSingleHandler(h.id, func(payload []byte) ([]byte, *RemoteErr) {
   507  		req := h.NewRequest()
   508  		_, err := req.UnmarshalMsg(payload)
   509  		if err != nil {
   510  			PutByteBuffer(payload)
   511  			r := RemoteErr(err.Error())
   512  			return nil, &r
   513  		}
   514  		resp, rerr := handle(req)
   515  		h.recycleReq(req)
   516  
   517  		if rerr != nil {
   518  			PutByteBuffer(payload)
   519  			return nil, rerr
   520  		}
   521  		payload, err = resp.MarshalMsg(payload[:0])
   522  		if !h.sharedResp {
   523  			h.PutResponse(resp)
   524  		}
   525  		if err != nil {
   526  			PutByteBuffer(payload)
   527  			r := RemoteErr(err.Error())
   528  			return nil, &r
   529  		}
   530  		return payload, nil
   531  	}, subroute...)
   532  }
   533  
   534  // Requester is able to send requests to a remote.
   535  type Requester interface {
   536  	Request(ctx context.Context, h HandlerID, req []byte) ([]byte, error)
   537  }
   538  
   539  // Call the remote with the request and return the response.
   540  // The response should be returned with PutResponse when no error.
   541  // If no deadline is set, a 1-minute deadline is added.
   542  func (h *SingleHandler[Req, Resp]) Call(ctx context.Context, c Requester, req Req) (resp Resp, err error) {
   543  	if c == nil {
   544  		if h.ignoreNilConn {
   545  			return resp, nil
   546  		}
   547  		return resp, ErrDisconnected
   548  	}
   549  	payload, err := req.MarshalMsg(GetByteBuffer()[:0])
   550  	if err != nil {
   551  		return resp, err
   552  	}
   553  	switch any(req).(type) {
   554  	case *MSS, *URLValues:
   555  		ctx = context.WithValue(ctx, TraceParamsKey{}, req)
   556  	case *NoPayload, *Bytes:
   557  		// do not need to trace nopayload and bytes payload
   558  	default:
   559  		ctx = context.WithValue(ctx, TraceParamsKey{}, fmt.Sprintf("type=%T", req))
   560  	}
   561  	if h.callReuseReq {
   562  		defer h.recycleReq(req)
   563  	}
   564  	res, err := c.Request(ctx, h.id, payload)
   565  	PutByteBuffer(payload)
   566  	if err != nil {
   567  		return resp, err
   568  	}
   569  	defer PutByteBuffer(res)
   570  	r := h.NewResponse()
   571  	_, err = r.UnmarshalMsg(res)
   572  	if err != nil {
   573  		h.PutResponse(r)
   574  		return resp, err
   575  	}
   576  	return r, err
   577  }
   578  
   579  // RemoteClient contains information about the caller.
   580  type RemoteClient struct {
   581  	Name string
   582  }
   583  
   584  type (
   585  	ctxCallerKey   = struct{}
   586  	ctxSubrouteKey = struct{}
   587  )
   588  
   589  // GetCaller returns caller information from contexts provided to handlers.
   590  func GetCaller(ctx context.Context) *RemoteClient {
   591  	val, _ := ctx.Value(ctxCallerKey{}).(*RemoteClient)
   592  	return val
   593  }
   594  
   595  // GetSubroute returns caller information from contexts provided to handlers.
   596  func GetSubroute(ctx context.Context) string {
   597  	val, _ := ctx.Value(ctxSubrouteKey{}).(string)
   598  	return val
   599  }
   600  
   601  func setCaller(ctx context.Context, cl *RemoteClient) context.Context {
   602  	return context.WithValue(ctx, ctxCallerKey{}, cl)
   603  }
   604  
   605  func setSubroute(ctx context.Context, s string) context.Context {
   606  	return context.WithValue(ctx, ctxSubrouteKey{}, s)
   607  }
   608  
   609  // StreamTypeHandler is a type safe handler for streaming requests.
   610  type StreamTypeHandler[Payload, Req, Resp RoundTripper] struct {
   611  	WithPayload bool
   612  
   613  	// Override the default capacities (1)
   614  	OutCapacity int
   615  
   616  	// Set to 0 if no input is expected.
   617  	// Will be 0 if newReq is nil.
   618  	InCapacity int
   619  
   620  	reqPool        sync.Pool
   621  	respPool       sync.Pool
   622  	id             HandlerID
   623  	newPayload     func() Payload
   624  	nilReq         Req
   625  	nilResp        Resp
   626  	sharedResponse bool
   627  }
   628  
   629  // NewStream creates a typed handler that can provide Marshal/Unmarshal.
   630  // Use Register to register a server handler.
   631  // Use Call to initiate a clientside call.
   632  // newPayload can be nil. In that case payloads will always be nil.
   633  // newReq can be nil. In that case no input stream is expected and the handler will be called with nil 'in' channel.
   634  func NewStream[Payload, Req, Resp RoundTripper](h HandlerID, newPayload func() Payload, newReq func() Req, newResp func() Resp) *StreamTypeHandler[Payload, Req, Resp] {
   635  	if newResp == nil {
   636  		panic("newResp missing in NewStream")
   637  	}
   638  
   639  	s := newStreamHandler[Payload, Req, Resp](h)
   640  	if newReq != nil {
   641  		s.reqPool.New = func() interface{} {
   642  			return newReq()
   643  		}
   644  	} else {
   645  		s.InCapacity = 0
   646  	}
   647  	s.respPool.New = func() interface{} {
   648  		return newResp()
   649  	}
   650  	s.newPayload = newPayload
   651  	s.WithPayload = newPayload != nil
   652  	return s
   653  }
   654  
   655  // WithSharedResponse indicates it is unsafe to reuse the response.
   656  // Typically this is used when the response sharing part of its data structure.
   657  func (h *StreamTypeHandler[Payload, Req, Resp]) WithSharedResponse() *StreamTypeHandler[Payload, Req, Resp] {
   658  	h.sharedResponse = true
   659  	return h
   660  }
   661  
   662  // NewPayload creates a new payload.
   663  func (h *StreamTypeHandler[Payload, Req, Resp]) NewPayload() Payload {
   664  	return h.newPayload()
   665  }
   666  
   667  // NewRequest creates a new request.
   668  // The struct may be reused, so caller should clear any fields.
   669  func (h *StreamTypeHandler[Payload, Req, Resp]) NewRequest() Req {
   670  	return h.reqPool.Get().(Req)
   671  }
   672  
   673  // PutRequest will accept a request for reuse.
   674  // These should be returned by the handler.
   675  func (h *StreamTypeHandler[Payload, Req, Resp]) PutRequest(r Req) {
   676  	if r != h.nilReq {
   677  		h.reqPool.Put(r)
   678  	}
   679  }
   680  
   681  // PutResponse will accept a response for reuse.
   682  // These should be returned by the caller.
   683  func (h *StreamTypeHandler[Payload, Req, Resp]) PutResponse(r Resp) {
   684  	if r != h.nilResp {
   685  		h.respPool.Put(r)
   686  	}
   687  }
   688  
   689  // NewResponse creates a new response.
   690  // Handlers can use this to create a reusable response.
   691  func (h *StreamTypeHandler[Payload, Req, Resp]) NewResponse() Resp {
   692  	return h.respPool.Get().(Resp)
   693  }
   694  
   695  func newStreamHandler[Payload, Req, Resp RoundTripper](h HandlerID) *StreamTypeHandler[Payload, Req, Resp] {
   696  	return &StreamTypeHandler[Payload, Req, Resp]{id: h, InCapacity: 1, OutCapacity: 1}
   697  }
   698  
   699  // Register a handler for two-way streaming with payload, input stream and output stream.
   700  // An optional subroute can be given. Multiple entries are joined with '/'.
   701  func (h *StreamTypeHandler[Payload, Req, Resp]) Register(m *Manager, handle func(ctx context.Context, p Payload, in <-chan Req, out chan<- Resp) *RemoteErr, subroute ...string) error {
   702  	return h.register(m, handle, subroute...)
   703  }
   704  
   705  // WithOutCapacity adjusts the output capacity from the handler perspective.
   706  // This must be done prior to registering the handler.
   707  func (h *StreamTypeHandler[Payload, Req, Resp]) WithOutCapacity(out int) *StreamTypeHandler[Payload, Req, Resp] {
   708  	h.OutCapacity = out
   709  	return h
   710  }
   711  
   712  // WithInCapacity adjusts the input capacity from the handler perspective.
   713  // This must be done prior to registering the handler.
   714  func (h *StreamTypeHandler[Payload, Req, Resp]) WithInCapacity(in int) *StreamTypeHandler[Payload, Req, Resp] {
   715  	h.InCapacity = in
   716  	return h
   717  }
   718  
   719  // RegisterNoInput a handler for one-way streaming with payload and output stream.
   720  // An optional subroute can be given. Multiple entries are joined with '/'.
   721  func (h *StreamTypeHandler[Payload, Req, Resp]) RegisterNoInput(m *Manager, handle func(ctx context.Context, p Payload, out chan<- Resp) *RemoteErr, subroute ...string) error {
   722  	h.InCapacity = 0
   723  	return h.register(m, func(ctx context.Context, p Payload, in <-chan Req, out chan<- Resp) *RemoteErr {
   724  		return handle(ctx, p, out)
   725  	}, subroute...)
   726  }
   727  
   728  // RegisterNoPayload a handler for one-way streaming with payload and output stream.
   729  // An optional subroute can be given. Multiple entries are joined with '/'.
   730  func (h *StreamTypeHandler[Payload, Req, Resp]) RegisterNoPayload(m *Manager, handle func(ctx context.Context, in <-chan Req, out chan<- Resp) *RemoteErr, subroute ...string) error {
   731  	h.WithPayload = false
   732  	return h.register(m, func(ctx context.Context, p Payload, in <-chan Req, out chan<- Resp) *RemoteErr {
   733  		return handle(ctx, in, out)
   734  	}, subroute...)
   735  }
   736  
   737  // Register a handler for two-way streaming with optional payload and input stream.
   738  func (h *StreamTypeHandler[Payload, Req, Resp]) register(m *Manager, handle func(ctx context.Context, p Payload, in <-chan Req, out chan<- Resp) *RemoteErr, subroute ...string) error {
   739  	return m.RegisterStreamingHandler(h.id, StreamHandler{
   740  		Handle: func(ctx context.Context, payload []byte, in <-chan []byte, out chan<- []byte) *RemoteErr {
   741  			var plT Payload
   742  			if h.WithPayload {
   743  				plT = h.NewPayload()
   744  				_, err := plT.UnmarshalMsg(payload)
   745  				PutByteBuffer(payload)
   746  				if err != nil {
   747  					r := RemoteErr(err.Error())
   748  					return &r
   749  				}
   750  			}
   751  
   752  			var inT chan Req
   753  			if h.InCapacity > 0 {
   754  				// Don't add extra buffering
   755  				inT = make(chan Req)
   756  				go func() {
   757  					defer xioutil.SafeClose(inT)
   758  					for {
   759  						select {
   760  						case <-ctx.Done():
   761  							return
   762  						case v, ok := <-in:
   763  							if !ok {
   764  								return
   765  							}
   766  							input := h.NewRequest()
   767  							_, err := input.UnmarshalMsg(v)
   768  							if err != nil {
   769  								logger.LogOnceIf(ctx, err, err.Error())
   770  							}
   771  							PutByteBuffer(v)
   772  							// Send input
   773  							select {
   774  							case <-ctx.Done():
   775  								return
   776  							case inT <- input:
   777  							}
   778  						}
   779  					}
   780  				}()
   781  			}
   782  			outT := make(chan Resp)
   783  			outDone := make(chan struct{})
   784  			go func() {
   785  				defer xioutil.SafeClose(outDone)
   786  				dropOutput := false
   787  				for v := range outT {
   788  					if dropOutput {
   789  						continue
   790  					}
   791  					dst := GetByteBuffer()
   792  					dst, err := v.MarshalMsg(dst[:0])
   793  					if err != nil {
   794  						logger.LogOnceIf(ctx, err, err.Error())
   795  					}
   796  					if !h.sharedResponse {
   797  						h.PutResponse(v)
   798  					}
   799  					select {
   800  					case <-ctx.Done():
   801  						dropOutput = true
   802  					case out <- dst:
   803  					}
   804  				}
   805  			}()
   806  			rErr := handle(ctx, plT, inT, outT)
   807  			xioutil.SafeClose(outT)
   808  			<-outDone
   809  			return rErr
   810  		}, OutCapacity: h.OutCapacity, InCapacity: h.InCapacity, Subroute: strings.Join(subroute, "/"),
   811  	})
   812  }
   813  
   814  // TypedStream is a stream with specific types.
   815  type TypedStream[Req, Resp RoundTripper] struct {
   816  	// responses from the remote server.
   817  	// Channel will be closed after error or when remote closes.
   818  	// responses *must* be read to either an error is returned or the channel is closed.
   819  	responses *Stream
   820  	newResp   func() Resp
   821  
   822  	// Requests sent to the server.
   823  	// If the handler is defined with 0 incoming capacity this will be nil.
   824  	// Channel *must* be closed to signal the end of the stream.
   825  	// If the request context is canceled, the stream will no longer process requests.
   826  	Requests chan<- Req
   827  }
   828  
   829  // Results returns the results from the remote server one by one.
   830  // If any error is returned by the callback, the stream will be canceled.
   831  // If the context is canceled, the stream will be canceled.
   832  func (s *TypedStream[Req, Resp]) Results(next func(resp Resp) error) (err error) {
   833  	return s.responses.Results(func(b []byte) error {
   834  		resp := s.newResp()
   835  		_, err := resp.UnmarshalMsg(b)
   836  		if err != nil {
   837  			return err
   838  		}
   839  		return next(resp)
   840  	})
   841  }
   842  
   843  // Streamer creates a stream.
   844  type Streamer interface {
   845  	NewStream(ctx context.Context, h HandlerID, payload []byte) (st *Stream, err error)
   846  }
   847  
   848  // Call the remove with the request and
   849  func (h *StreamTypeHandler[Payload, Req, Resp]) Call(ctx context.Context, c Streamer, payload Payload) (st *TypedStream[Req, Resp], err error) {
   850  	if c == nil {
   851  		return nil, ErrDisconnected
   852  	}
   853  	var payloadB []byte
   854  	if h.WithPayload {
   855  		var err error
   856  		payloadB, err = payload.MarshalMsg(GetByteBuffer()[:0])
   857  		if err != nil {
   858  			return nil, err
   859  		}
   860  	}
   861  	stream, err := c.NewStream(ctx, h.id, payloadB)
   862  	PutByteBuffer(payloadB)
   863  	if err != nil {
   864  		return nil, err
   865  	}
   866  
   867  	// respT := make(chan TypedResponse[Resp])
   868  	var reqT chan Req
   869  	if h.InCapacity > 0 {
   870  		reqT = make(chan Req)
   871  		// Request handler
   872  		if stream.Requests == nil {
   873  			return nil, fmt.Errorf("internal error: stream request channel nil")
   874  		}
   875  		go func() {
   876  			defer xioutil.SafeClose(stream.Requests)
   877  			for req := range reqT {
   878  				b, err := req.MarshalMsg(GetByteBuffer()[:0])
   879  				if err != nil {
   880  					logger.LogOnceIf(ctx, err, err.Error())
   881  				}
   882  				h.PutRequest(req)
   883  				stream.Requests <- b
   884  			}
   885  		}()
   886  	} else if stream.Requests != nil {
   887  		xioutil.SafeClose(stream.Requests)
   888  	}
   889  
   890  	return &TypedStream[Req, Resp]{responses: stream, newResp: h.NewResponse, Requests: reqT}, nil
   891  }