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