code.vegaprotocol.io/vega@v0.79.0/core/execution/common/expiring_orders.go (about)

     1  // Copyright (C) 2023 Gobalsky Labs Limited
     2  //
     3  // This program is free software: you can redistribute it and/or modify
     4  // it under the terms of the GNU Affero General Public License as
     5  // published by the Free Software Foundation, either version 3 of the
     6  // License, or (at your option) any later version.
     7  //
     8  // This program is distributed in the hope that it will be useful,
     9  // but WITHOUT ANY WARRANTY; without even the implied warranty of
    10  // MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
    11  // GNU Affero General Public License for more details.
    12  //
    13  // You should have received a copy of the GNU Affero General Public License
    14  // along with this program.  If not, see <http://www.gnu.org/licenses/>.
    15  
    16  package common
    17  
    18  import (
    19  	"github.com/google/btree"
    20  )
    21  
    22  type ExpiringOrders struct {
    23  	orders *btree.BTree
    24  }
    25  
    26  type ordersAtTS struct {
    27  	ts int64
    28  	// order IDs
    29  	orders []string
    30  }
    31  
    32  func (a *ordersAtTS) Less(b btree.Item) bool {
    33  	return a.ts < b.(*ordersAtTS).ts
    34  }
    35  
    36  func NewExpiringOrders() *ExpiringOrders {
    37  	return &ExpiringOrders{
    38  		orders: btree.New(2),
    39  	}
    40  }
    41  
    42  func (a *ExpiringOrders) GetExpiryingOrderCount() int {
    43  	result := a.orders.Len()
    44  	return result
    45  }
    46  
    47  func (a *ExpiringOrders) Insert(
    48  	orderID string, ts int64,
    49  ) {
    50  	item := &ordersAtTS{ts: ts}
    51  	if item := a.orders.Get(item); item != nil {
    52  		item.(*ordersAtTS).orders = append(item.(*ordersAtTS).orders, orderID)
    53  		return
    54  	}
    55  	item.orders = []string{orderID}
    56  	a.orders.ReplaceOrInsert(item)
    57  }
    58  
    59  func (a *ExpiringOrders) RemoveOrder(expiresAt int64, orderID string) bool {
    60  	item := &ordersAtTS{ts: expiresAt}
    61  	if item := a.orders.Get(item); item != nil {
    62  		oat := item.(*ordersAtTS)
    63  		for i := 0; i < len(oat.orders); i++ {
    64  			if oat.orders[i] == orderID {
    65  				oat.orders = oat.orders[:i+copy(oat.orders[i:], oat.orders[i+1:])]
    66  
    67  				// if the slice is empty, remove the parent container
    68  				if len(oat.orders) == 0 {
    69  					a.orders.Delete(item)
    70  				}
    71  				return true
    72  			}
    73  		}
    74  	}
    75  	return false
    76  }
    77  
    78  func (a *ExpiringOrders) Expire(ts int64) []string {
    79  	if a.orders.Len() == 0 {
    80  		return nil
    81  	}
    82  	orders := []string{}
    83  	toDelete := []int64{}
    84  	item := &ordersAtTS{ts: ts + 1}
    85  	a.orders.AscendLessThan(item, func(i btree.Item) bool {
    86  		if ts < i.(*ordersAtTS).ts {
    87  			return false
    88  		}
    89  		orders = append(orders, i.(*ordersAtTS).orders...)
    90  		toDelete = append(toDelete, i.(*ordersAtTS).ts)
    91  		return true
    92  	})
    93  
    94  	for _, v := range toDelete {
    95  		item.ts = v
    96  		a.orders.Delete(item)
    97  	}
    98  
    99  	return orders
   100  }