go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/logdog/client/butler/bundler/data.go (about)

     1  // Copyright 2015 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package bundler
    16  
    17  import (
    18  	"sync"
    19  	"time"
    20  
    21  	"go.chromium.org/luci/common/data/chunkstream"
    22  )
    23  
    24  type dataPoolRegistry struct {
    25  	sync.Mutex
    26  
    27  	// pools is a pool of Data instances. It is keyed on buffer size.
    28  	pools map[int]*dataPool
    29  }
    30  
    31  // globaldataPoolRegistry is the default data pool to use by the stream
    32  // package.
    33  var globalDataPoolRegistry = dataPoolRegistry{}
    34  
    35  func (r *dataPoolRegistry) getPool(s int) *dataPool {
    36  	r.Lock()
    37  	defer r.Unlock()
    38  
    39  	if r.pools == nil {
    40  		r.pools = map[int]*dataPool{}
    41  	}
    42  
    43  	pool := r.pools[s]
    44  	if pool == nil {
    45  		pool = newPool(s)
    46  		r.pools[s] = pool
    47  	}
    48  	return pool
    49  }
    50  
    51  type dataPool struct {
    52  	sync.Pool
    53  	size int
    54  }
    55  
    56  func newPool(size int) *dataPool {
    57  	p := dataPool{
    58  		size: size,
    59  	}
    60  	p.New = p.newData
    61  	return &p
    62  }
    63  
    64  func (p *dataPool) newData() any {
    65  	return &streamData{
    66  		buffer:      make([]byte, p.size),
    67  		releaseFunc: p.release,
    68  	}
    69  }
    70  
    71  func (p *dataPool) getData() Data {
    72  	d := p.Get().(*streamData)
    73  	d.reset()
    74  	return d
    75  }
    76  
    77  func (p *dataPool) release(d *streamData) {
    78  	p.Put(d)
    79  }
    80  
    81  // Data is a reusable data buffer that is used by Stream instances to ingest
    82  // data.
    83  //
    84  // Data is initially an empty buffer. Once data is loaded into it, the buffer is
    85  // resized to the bound data and a timestamp is attached via Bind.
    86  type Data interface {
    87  	chunkstream.Chunk
    88  
    89  	// Bind resizes the Chunk buffer and records a timestamp to associate with the
    90  	// data chunk.
    91  	Bind(int, time.Time) Data
    92  
    93  	// Timestamp returns the bound timestamp. This will be zero if no timestamp
    94  	// has been bound.
    95  	Timestamp() time.Time
    96  }
    97  
    98  // streamData is an implementation of the Chunk interface for Bundler chunks.
    99  //
   100  // It includes the ability to bind to a size/timestamp.
   101  type streamData struct {
   102  	buffer []byte
   103  	ts     time.Time
   104  
   105  	releaseFunc func(*streamData)
   106  }
   107  
   108  var _ Data = (*streamData)(nil)
   109  
   110  func (d *streamData) reset() {
   111  	d.buffer = d.buffer[:cap(d.buffer)]
   112  }
   113  
   114  func (d *streamData) Bytes() []byte {
   115  	return d.buffer
   116  }
   117  
   118  func (d *streamData) Len() int {
   119  	return len(d.buffer)
   120  }
   121  
   122  func (d *streamData) Bind(amount int, ts time.Time) Data {
   123  	d.buffer = d.buffer[:amount]
   124  	d.ts = ts
   125  	return d
   126  }
   127  
   128  func (d *streamData) Timestamp() time.Time {
   129  	return d.ts
   130  }
   131  
   132  func (d *streamData) Release() {
   133  	if d.releaseFunc != nil {
   134  		d.releaseFunc(d)
   135  	}
   136  }