github.com/m3db/m3@v1.5.1-0.20231129193456-75a402aa583b/src/dbnode/client/fetch_batch.go (about)

     1  // Copyright (c) 2016 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package client
    22  
    23  import (
    24  	"github.com/m3db/m3/src/dbnode/generated/thrift/rpc"
    25  	"github.com/m3db/m3/src/x/checked"
    26  	"github.com/m3db/m3/src/x/pool"
    27  )
    28  
    29  type fetchBatchOp struct {
    30  	checked.RefCount
    31  	request           rpc.FetchBatchRawRequest
    32  	requestV2Elements []rpc.FetchBatchRawV2RequestElement
    33  	completionFns     []completionFn
    34  	finalizer         fetchBatchOpFinalizer
    35  }
    36  
    37  func (f *fetchBatchOp) reset() {
    38  	f.IncWrites()
    39  	f.request.RangeStart = 0
    40  	f.request.RangeEnd = 0
    41  	f.request.NameSpace = nil
    42  	for i := range f.request.Ids {
    43  		f.request.Ids[i] = nil
    44  	}
    45  	f.request.Ids = f.request.Ids[:0]
    46  	for i := range f.completionFns {
    47  		f.completionFns[i] = nil
    48  	}
    49  	f.completionFns = f.completionFns[:0]
    50  
    51  	for i := range f.requestV2Elements {
    52  		f.requestV2Elements[i].NameSpace = 0
    53  		f.requestV2Elements[i].RangeStart = 0
    54  		f.requestV2Elements[i].RangeEnd = 0
    55  		f.requestV2Elements[i].ID = nil
    56  		f.requestV2Elements[i].RangeTimeType = 0
    57  	}
    58  	f.requestV2Elements = f.requestV2Elements[:0]
    59  
    60  	f.DecWrites()
    61  }
    62  
    63  func (f *fetchBatchOp) append(namespace, id []byte, completionFn completionFn) {
    64  	f.IncWrites()
    65  	f.request.NameSpace = namespace
    66  	f.request.Ids = append(f.request.Ids, id)
    67  	f.requestV2Elements = append(f.requestV2Elements, rpc.FetchBatchRawV2RequestElement{
    68  		// NameSpace filled in by the host queue later.
    69  		RangeStart:    f.request.RangeStart,
    70  		RangeEnd:      f.request.RangeEnd,
    71  		ID:            id,
    72  		RangeTimeType: f.request.RangeTimeType,
    73  	})
    74  	f.completionFns = append(f.completionFns, completionFn)
    75  	f.DecWrites()
    76  }
    77  
    78  func (f *fetchBatchOp) Size() int {
    79  	f.IncReads()
    80  	value := len(f.request.Ids)
    81  	if value == 0 {
    82  		value = len(f.requestV2Elements)
    83  	}
    84  	f.DecReads()
    85  	return value
    86  }
    87  
    88  func (f *fetchBatchOp) CompletionFn() completionFn {
    89  	return f.completeAll
    90  }
    91  
    92  func (f *fetchBatchOp) completeAll(result interface{}, err error) {
    93  	for idx := range f.completionFns {
    94  		f.completionFns[idx](result, err)
    95  	}
    96  }
    97  
    98  func (f *fetchBatchOp) complete(idx int, result interface{}, err error) {
    99  	f.IncReads()
   100  	fn := f.completionFns[idx]
   101  	f.DecReads()
   102  	fn(result, err)
   103  }
   104  
   105  type fetchBatchOpFinalizer struct {
   106  	ref  *fetchBatchOp
   107  	pool *fetchBatchOpPool
   108  }
   109  
   110  func (f fetchBatchOpFinalizer) OnFinalize() {
   111  	f.pool.Put(f.ref)
   112  }
   113  
   114  type fetchBatchOpPool struct {
   115  	pool     pool.ObjectPool
   116  	capacity int
   117  }
   118  
   119  func newFetchBatchOpPool(opts pool.ObjectPoolOptions, capacity int) *fetchBatchOpPool {
   120  	p := pool.NewObjectPool(opts)
   121  	return &fetchBatchOpPool{pool: p, capacity: capacity}
   122  }
   123  
   124  func (p *fetchBatchOpPool) Init() {
   125  	p.pool.Init(func() interface{} {
   126  		f := &fetchBatchOp{}
   127  		f.request.Ids = make([][]byte, 0, p.capacity)
   128  		f.completionFns = make([]completionFn, 0, p.capacity)
   129  		f.finalizer.ref = f
   130  		f.finalizer.pool = p
   131  		f.SetOnFinalize(&f.finalizer)
   132  
   133  		f.IncRef()
   134  		f.reset()
   135  		f.DecRef()
   136  		return f
   137  	})
   138  }
   139  
   140  func (p *fetchBatchOpPool) Get() *fetchBatchOp {
   141  	return p.pool.Get().(*fetchBatchOp)
   142  }
   143  
   144  func (p *fetchBatchOpPool) Put(f *fetchBatchOp) {
   145  	f.IncRef()
   146  	if cap(f.request.Ids) != p.capacity || cap(f.completionFns) != p.capacity {
   147  		// Grew outside capacity, do not return to pool
   148  		f.DecRef()
   149  		return
   150  	}
   151  	f.reset()
   152  	f.DecRef()
   153  	p.pool.Put(f)
   154  }