code.vegaprotocol.io/vega@v0.79.0/core/coreapi/services/votes.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 services
    17  
    18  import (
    19  	"context"
    20  	"errors"
    21  	"sync"
    22  
    23  	"code.vegaprotocol.io/vega/core/events"
    24  	"code.vegaprotocol.io/vega/core/subscribers"
    25  	vegapb "code.vegaprotocol.io/vega/protos/vega"
    26  )
    27  
    28  var ErrMissingProposalOrPartyFilter = errors.New("missing proposal or party filter")
    29  
    30  type voteE interface {
    31  	events.Event
    32  	Vote() vegapb.Vote
    33  }
    34  
    35  type Votes struct {
    36  	*subscribers.Base
    37  	ctx context.Context
    38  
    39  	mu sync.RWMutex
    40  	// map of proposal id -> vote id -> vote
    41  	votes map[string]map[string]vegapb.Vote
    42  	// map of proposer -> set of vote id
    43  	votesPerParty map[string]map[string]struct{}
    44  	ch            chan vegapb.Vote
    45  }
    46  
    47  func NewVotes(ctx context.Context) (votes *Votes) {
    48  	defer func() { go votes.consume() }()
    49  	return &Votes{
    50  		Base:          subscribers.NewBase(ctx, 1000, true),
    51  		ctx:           ctx,
    52  		votes:         map[string]map[string]vegapb.Vote{},
    53  		votesPerParty: map[string]map[string]struct{}{},
    54  		ch:            make(chan vegapb.Vote, 100),
    55  	}
    56  }
    57  
    58  func (v *Votes) consume() {
    59  	defer func() { close(v.ch) }()
    60  	for {
    61  		select {
    62  		case <-v.Closed():
    63  			return
    64  		case vote, ok := <-v.ch:
    65  			if !ok {
    66  				// cleanup base
    67  				v.Halt()
    68  				// channel is closed
    69  				return
    70  			}
    71  			v.mu.Lock()
    72  			// first add to the proposals maps
    73  			votes, ok := v.votes[vote.ProposalId]
    74  			if !ok {
    75  				votes = map[string]vegapb.Vote{}
    76  				v.votes[vote.ProposalId] = votes
    77  			}
    78  			votes[vote.PartyId] = vote
    79  
    80  			// next to the party
    81  			partyVotes, ok := v.votesPerParty[vote.PartyId]
    82  			if !ok {
    83  				partyVotes = map[string]struct{}{}
    84  				v.votesPerParty[vote.PartyId] = partyVotes
    85  			}
    86  			partyVotes[vote.ProposalId] = struct{}{}
    87  			v.mu.Unlock()
    88  		}
    89  	}
    90  }
    91  
    92  func (v *Votes) Push(evts ...events.Event) {
    93  	for _, e := range evts {
    94  		if ae, ok := e.(voteE); ok {
    95  			v.ch <- ae.Vote()
    96  		}
    97  	}
    98  }
    99  
   100  func (v *Votes) List(proposal, party string) ([]*vegapb.Vote, error) {
   101  	v.mu.RLock()
   102  	defer v.mu.RUnlock()
   103  	if len(proposal) > 0 && len(party) > 0 {
   104  		return v.getVotesPerProposalAndParty(proposal, party), nil
   105  	} else if len(party) > 0 {
   106  		return v.getPartyVotes(party), nil
   107  	} else if len(proposal) > 0 {
   108  		return v.getProposalVotes(proposal), nil
   109  	}
   110  	return nil, ErrMissingProposalOrPartyFilter
   111  }
   112  
   113  func (v *Votes) getVotesPerProposalAndParty(proposal, party string) []*vegapb.Vote {
   114  	out := []*vegapb.Vote{}
   115  	propVotes, ok := v.votes[proposal]
   116  	if !ok {
   117  		return out
   118  	}
   119  
   120  	vote, ok := propVotes[party]
   121  	if ok {
   122  		out = append(out, &vote)
   123  	}
   124  
   125  	return out
   126  }
   127  
   128  func (v *Votes) getPartyVotes(party string) []*vegapb.Vote {
   129  	partyVotes, ok := v.votesPerParty[party]
   130  	if !ok {
   131  		return nil
   132  	}
   133  
   134  	out := make([]*vegapb.Vote, 0, len(partyVotes))
   135  	for k := range partyVotes {
   136  		vote := v.votes[k][party]
   137  		out = append(out, &vote)
   138  	}
   139  	return out
   140  }
   141  
   142  func (v *Votes) getProposalVotes(proposal string) []*vegapb.Vote {
   143  	proposalVotes, ok := v.votes[proposal]
   144  	if !ok {
   145  		return nil
   146  	}
   147  
   148  	out := make([]*vegapb.Vote, 0, len(proposalVotes))
   149  	for _, v := range proposalVotes {
   150  		v := v
   151  		out = append(out, &v)
   152  	}
   153  	return out
   154  }
   155  
   156  func (v *Votes) Types() []events.Type {
   157  	return []events.Type{
   158  		events.VoteEvent,
   159  	}
   160  }