github.com/MetalBlockchain/metalgo@v1.11.9/chains/atomic/memory.go (about)

     1  // Copyright (C) 2019-2024, Ava Labs, Inc. All rights reserved.
     2  // See the file LICENSE for licensing terms.
     3  
     4  package atomic
     5  
     6  import (
     7  	"bytes"
     8  	"sync"
     9  
    10  	"github.com/MetalBlockchain/metalgo/database"
    11  	"github.com/MetalBlockchain/metalgo/database/prefixdb"
    12  	"github.com/MetalBlockchain/metalgo/ids"
    13  	"github.com/MetalBlockchain/metalgo/utils/hashing"
    14  )
    15  
    16  type rcLock struct {
    17  	lock  sync.Mutex
    18  	count int
    19  }
    20  
    21  // Memory is used to set up a bidirectional communication channel between a pair
    22  // of chains.
    23  //
    24  // For any such pair, we compute a hash of the ordered pair of IDs to use as a
    25  // prefix DB that can be shared across the two chains. On top of the prefix DB
    26  // shared among two chains, we use constant prefixes to determine the
    27  // inbound/outbound and value/index database assignments.
    28  type Memory struct {
    29  	lock  sync.Mutex
    30  	locks map[ids.ID]*rcLock
    31  	db    database.Database
    32  }
    33  
    34  func NewMemory(db database.Database) *Memory {
    35  	return &Memory{
    36  		locks: make(map[ids.ID]*rcLock),
    37  		db:    db,
    38  	}
    39  }
    40  
    41  func (m *Memory) NewSharedMemory(chainID ids.ID) SharedMemory {
    42  	return &sharedMemory{
    43  		m:           m,
    44  		thisChainID: chainID,
    45  	}
    46  }
    47  
    48  // GetSharedDatabase returns a new locked prefix db on top of an existing
    49  // database
    50  //
    51  // Invariant: ReleaseSharedDatabase must be called after to free the database
    52  // associated with [sharedID]
    53  func (m *Memory) GetSharedDatabase(db database.Database, sharedID ids.ID) database.Database {
    54  	lock := m.makeLock(sharedID)
    55  	lock.Lock()
    56  	return prefixdb.NewNested(sharedID[:], db)
    57  }
    58  
    59  // ReleaseSharedDatabase unlocks the provided DB
    60  //
    61  // Note: ReleaseSharedDatabase must be called only after a corresponding call to
    62  // GetSharedDatabase. If ReleaseSharedDatabase is called without a corresponding
    63  // one-to-one call with GetSharedDatabase, it will panic.
    64  func (m *Memory) ReleaseSharedDatabase(sharedID ids.ID) {
    65  	lock := m.releaseLock(sharedID)
    66  	lock.Unlock()
    67  }
    68  
    69  // makeLock returns the lock associated with [sharedID], or creates a new one if
    70  // it doesn't exist yet, and increments the reference count.
    71  func (m *Memory) makeLock(sharedID ids.ID) *sync.Mutex {
    72  	m.lock.Lock()
    73  	defer m.lock.Unlock()
    74  
    75  	rc, exists := m.locks[sharedID]
    76  	if !exists {
    77  		rc = &rcLock{}
    78  		m.locks[sharedID] = rc
    79  	}
    80  	rc.count++
    81  	return &rc.lock
    82  }
    83  
    84  // releaseLock returns the lock associated with [sharedID] and decrements its
    85  // reference count. If this brings the count to 0, it will remove the lock from
    86  // the internal map of locks. If there is no lock associated with [sharedID],
    87  // releaseLock will panic.
    88  func (m *Memory) releaseLock(sharedID ids.ID) *sync.Mutex {
    89  	m.lock.Lock()
    90  	defer m.lock.Unlock()
    91  
    92  	rc, exists := m.locks[sharedID]
    93  	if !exists {
    94  		panic("attempting to free an unknown lock")
    95  	}
    96  	rc.count--
    97  	if rc.count == 0 {
    98  		delete(m.locks, sharedID)
    99  	}
   100  	return &rc.lock
   101  }
   102  
   103  // sharedID calculates the ID of the shared memory space
   104  func sharedID(id1, id2 ids.ID) ids.ID {
   105  	// Swap IDs locally to ensure id1 <= id2.
   106  	if bytes.Compare(id1[:], id2[:]) == 1 {
   107  		id1, id2 = id2, id1
   108  	}
   109  
   110  	combinedBytes, err := Codec.Marshal(CodecVersion, [2]ids.ID{id1, id2})
   111  	if err != nil {
   112  		panic(err)
   113  	}
   114  	return hashing.ComputeHash256Array(combinedBytes)
   115  }