github.com/m3db/m3@v1.5.0/src/msg/producer/ref_counted.go (about)

     1  // Copyright (c) 2018 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 producer
    22  
    23  import (
    24  	"sync"
    25  
    26  	"go.uber.org/atomic"
    27  )
    28  
    29  // OnFinalizeFn will be called when the message is being finalized.
    30  type OnFinalizeFn func(rm *RefCountedMessage)
    31  
    32  // RefCountedMessage is a reference counted message.
    33  type RefCountedMessage struct {
    34  	mu sync.RWMutex
    35  	Message
    36  
    37  	size         uint64
    38  	onFinalizeFn OnFinalizeFn
    39  
    40  	// RefCountedMessage must not be copied by value due to RWMutex,
    41  	// safe to store values here and not just pointers
    42  	refCount            atomic.Int32
    43  	isDroppedOrConsumed atomic.Bool
    44  }
    45  
    46  // NewRefCountedMessage creates RefCountedMessage.
    47  func NewRefCountedMessage(m Message, fn OnFinalizeFn) *RefCountedMessage {
    48  	return &RefCountedMessage{
    49  		Message:      m,
    50  		size:         uint64(m.Size()),
    51  		onFinalizeFn: fn,
    52  	}
    53  }
    54  
    55  // Accept returns true if the message can be accepted by the filter.
    56  func (rm *RefCountedMessage) Accept(fn FilterFunc) bool {
    57  	return fn(rm.Message)
    58  }
    59  
    60  // IncRef increments the ref count.
    61  func (rm *RefCountedMessage) IncRef() {
    62  	rm.refCount.Inc()
    63  }
    64  
    65  // DecRef decrements the ref count. If the reference count became zero after
    66  // the call, the message will be finalized as consumed.
    67  func (rm *RefCountedMessage) DecRef() {
    68  	rc := rm.refCount.Dec()
    69  	if rc == 0 {
    70  		rm.finalize(Consumed)
    71  	}
    72  	if rc < 0 {
    73  		panic("invalid ref count")
    74  	}
    75  }
    76  
    77  // IncReads increments the reads count.
    78  func (rm *RefCountedMessage) IncReads() {
    79  	rm.mu.RLock()
    80  }
    81  
    82  // DecReads decrements the reads count.
    83  func (rm *RefCountedMessage) DecReads() {
    84  	rm.mu.RUnlock()
    85  }
    86  
    87  // NumRef returns the number of references remaining.
    88  func (rm *RefCountedMessage) NumRef() int32 {
    89  	return rm.refCount.Load()
    90  }
    91  
    92  // Size returns the size of the message.
    93  func (rm *RefCountedMessage) Size() uint64 {
    94  	return rm.size
    95  }
    96  
    97  // Drop drops the message without waiting for it to be consumed.
    98  func (rm *RefCountedMessage) Drop() bool {
    99  	return rm.finalize(Dropped)
   100  }
   101  
   102  // IsDroppedOrConsumed returns true if the message has been dropped or consumed.
   103  func (rm *RefCountedMessage) IsDroppedOrConsumed() bool {
   104  	return rm.isDroppedOrConsumed.Load()
   105  }
   106  
   107  func (rm *RefCountedMessage) finalize(r FinalizeReason) bool {
   108  	// NB: This lock prevents the message from being finalized when its still
   109  	// being read.
   110  	rm.mu.Lock()
   111  	if rm.isDroppedOrConsumed.Load() {
   112  		rm.mu.Unlock()
   113  		return false
   114  	}
   115  	rm.isDroppedOrConsumed.Store(true)
   116  	rm.mu.Unlock()
   117  	if rm.onFinalizeFn != nil {
   118  		rm.onFinalizeFn(rm)
   119  	}
   120  	rm.Message.Finalize(r)
   121  	return true
   122  }