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 }