github.com/juju/juju@v0.0.0-20240430160146-1752b71fcf00/state/querytracker.go (about)

     1  // Copyright 2020 Canonical Ltd.
     2  // Licensed under the AGPLv3, see LICENCE file for details.
     3  
     4  package state
     5  
     6  import (
     7  	"runtime/debug"
     8  	"strings"
     9  	"sync"
    10  )
    11  
    12  // QueryTracker provides a way for tests to determine how many
    13  // database queries have been made, and who made them.
    14  type QueryTracker interface {
    15  	Reset()
    16  	ListQueries() []QueryDetails
    17  	ReadCount() int
    18  	// TODO: implement the write tracking.
    19  	// WriteCount() int
    20  }
    21  
    22  // QueryDetails is a POD type recording the database query
    23  // and who made it.
    24  type QueryDetails struct {
    25  	Type           string // read or write
    26  	CollectionName string
    27  	Query          interface{}
    28  	Traceback      string
    29  }
    30  
    31  // TrackQueries allows tests to turn on a mechanism to count and
    32  // track the database queries made. The query is counted if
    33  // the method specified is found in the traceback. It is currently
    34  // just using a string.Contains check. We may want to extend this
    35  // functionality at some state to use regex, or support multiple
    36  // matches.
    37  func (s *State) TrackQueries(method string) QueryTracker {
    38  	tracker := &queryTracker{method: method}
    39  	s.database.(*database).setTracker(tracker)
    40  	return tracker
    41  }
    42  
    43  type queryTracker struct {
    44  	method  string
    45  	mu      sync.Mutex
    46  	queries []QueryDetails
    47  }
    48  
    49  // Reset clears out all the current reads and writes.
    50  func (q *queryTracker) Reset() {
    51  	q.mu.Lock()
    52  	defer q.mu.Unlock()
    53  
    54  	q.queries = nil
    55  }
    56  
    57  // ListQueries returns the list of all queries that have been
    58  // done since start or reset.
    59  func (q *queryTracker) ListQueries() []QueryDetails {
    60  	q.mu.Lock()
    61  	defer q.mu.Unlock()
    62  
    63  	return q.queries
    64  }
    65  
    66  // ReadCount returns the number of read queries that have been
    67  // done since start or reset.
    68  func (q *queryTracker) ReadCount() int {
    69  	q.mu.Lock()
    70  	defer q.mu.Unlock()
    71  
    72  	count := 0
    73  	for _, query := range q.queries {
    74  		if query.Type == "read" {
    75  			count++
    76  		}
    77  	}
    78  	return count
    79  }
    80  
    81  // WriteCount returns the number of write queries that have been
    82  // done since start or reset.
    83  func (q *queryTracker) WriteCount() int {
    84  	q.mu.Lock()
    85  	defer q.mu.Unlock()
    86  
    87  	count := 0
    88  	for _, query := range q.queries {
    89  		if query.Type == "write" {
    90  			count++
    91  		}
    92  	}
    93  	return count
    94  }
    95  
    96  // TrackRead records the read query against the collection specified
    97  // and where the call came from.
    98  func (q *queryTracker) TrackRead(collectionName string, query interface{}) {
    99  	q.mu.Lock()
   100  	defer q.mu.Unlock()
   101  
   102  	traceback := string(debug.Stack())
   103  	if strings.Contains(traceback, q.method) {
   104  		q.queries = append(q.queries, QueryDetails{
   105  			Type:           "read",
   106  			CollectionName: collectionName,
   107  			Query:          query,
   108  			Traceback:      traceback,
   109  		})
   110  	}
   111  }