github.com/cloudwego/kitex@v0.9.0/pkg/utils/ring.go (about)

     1  /*
     2   * Copyright 2021 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package utils
    18  
    19  import (
    20  	"errors"
    21  	"runtime"
    22  
    23  	goid "github.com/choleraehyq/pid"
    24  )
    25  
    26  // ErrRingFull means the ring is full.
    27  var ErrRingFull = errors.New("ring is full")
    28  
    29  // NewRing creates a ringbuffer with fixed size.
    30  func NewRing(size int) *Ring {
    31  	if size <= 0 {
    32  		panic("new ring with size <=0")
    33  	}
    34  	r := &Ring{}
    35  
    36  	// check len(rings):
    37  	// 1. len = P/4
    38  	// 2. len <= size
    39  	numP := runtime.GOMAXPROCS(0)
    40  	r.length = (numP + 3) / 4
    41  	if r.length > size {
    42  		r.length = size
    43  	}
    44  
    45  	// make rings
    46  	per := size / r.length
    47  	r.rings = make([]*ring, r.length)
    48  	for i := 0; i < r.length-1; i++ {
    49  		r.rings[i] = newRing(per)
    50  	}
    51  	r.rings[r.length-1] = newRing(size - per*(r.length-1))
    52  	return r
    53  }
    54  
    55  // Ring implements a fixed size hash list to manage data
    56  type Ring struct {
    57  	length int
    58  	rings  []*ring
    59  }
    60  
    61  // Push appends item to the ring.
    62  func (r *Ring) Push(obj interface{}) error {
    63  	if r.length == 1 {
    64  		return r.rings[0].Push(obj)
    65  	}
    66  
    67  	idx := goid.GetPid() % r.length
    68  	for i := 0; i < r.length; i, idx = i+1, (idx+1)%r.length {
    69  		err := r.rings[idx].Push(obj)
    70  		if err == nil {
    71  			return nil
    72  		}
    73  	}
    74  	return ErrRingFull
    75  }
    76  
    77  // Pop returns the last item and removes it from the ring.
    78  func (r *Ring) Pop() interface{} {
    79  	if r.length == 1 {
    80  		return r.rings[0].Pop()
    81  	}
    82  
    83  	idx := goid.GetPid() % r.length
    84  	for i := 0; i < r.length; i, idx = i+1, (idx+1)%r.length {
    85  		obj := r.rings[idx].Pop()
    86  		if obj != nil {
    87  			return obj
    88  		}
    89  	}
    90  	return nil
    91  }
    92  
    93  // Dump dumps the data in the ring.
    94  func (r *Ring) Dump() interface{} {
    95  	m := &ringDump{}
    96  	dumpList := make([]*ringDump, 0, r.length)
    97  	idx := goid.GetPid() % r.length
    98  	for i := 0; i < r.length; i, idx = i+1, (idx+1)%r.length {
    99  		curDump := &ringDump{}
   100  		r.rings[idx].Dump(curDump)
   101  		dumpList = append(dumpList, curDump)
   102  		m.Cap += curDump.Cap
   103  		m.Len += curDump.Len
   104  	}
   105  	m.Array = make([]interface{}, 0, m.Len)
   106  	for _, shardData := range dumpList {
   107  		for i := 0; i < shardData.Len; i++ {
   108  			m.Array = append(m.Array, shardData.Array[i])
   109  		}
   110  	}
   111  	return m
   112  }