github.com/siglens/siglens@v0.0.0-20240328180423-f7ce9ae441ed/pkg/segment/results/mresults/tsid/tsidtracker.go (about)

     1  /*
     2  Copyright 2023.
     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 tsidtracker
    18  
    19  import (
    20  	"sort"
    21  
    22  	"github.com/valyala/bytebufferpool"
    23  )
    24  
    25  var TAG_VALUE_DELIMITER_BYTE = []byte("`")
    26  
    27  var TAG_VALUE_DELIMITER_STR = ("`")
    28  
    29  /*
    30  Holder struct to track all matched TSIDs
    31  */
    32  type AllMatchedTSIDs struct {
    33  	allTSIDs map[uint64]*bytebufferpool.ByteBuffer // raw tsids that are currently being tracked
    34  	first    bool
    35  }
    36  
    37  /*
    38  This function should initialize a TSID tracker
    39  */
    40  func InitTSIDTracker(numTagFilters int) (*AllMatchedTSIDs, error) {
    41  
    42  	return &AllMatchedTSIDs{
    43  		allTSIDs: make(map[uint64]*bytebufferpool.ByteBuffer, 0),
    44  		first:    true,
    45  	}, nil
    46  }
    47  
    48  // If first time, add all tsids to map
    49  // Else, intersect with existing tsids
    50  func (tr *AllMatchedTSIDs) BulkAdd(rawTagValueToTSIDs map[string]map[uint64]struct{}) error {
    51  	if tr.first {
    52  		for tagValue, tsids := range rawTagValueToTSIDs {
    53  			for id := range tsids {
    54  				buff := bytebufferpool.Get()
    55  				_, err := buff.WriteString(tagValue)
    56  				if err != nil {
    57  					return err
    58  				}
    59  				_, err = buff.Write(TAG_VALUE_DELIMITER_BYTE)
    60  				if err != nil {
    61  					return err
    62  				}
    63  				tr.allTSIDs[id] = buff
    64  			}
    65  		}
    66  	} else {
    67  		valid := 0
    68  		for ts, tsidInfo := range tr.allTSIDs {
    69  			shouldKeep := false
    70  			for tagValue, tsids := range rawTagValueToTSIDs {
    71  				if _, ok := tsids[ts]; ok {
    72  					shouldKeep = true
    73  					valid++
    74  
    75  					_, err := tsidInfo.WriteString(tagValue)
    76  					if err != nil {
    77  						return err
    78  					}
    79  					_, err = tsidInfo.Write(TAG_VALUE_DELIMITER_BYTE)
    80  					if err != nil {
    81  						return err
    82  					}
    83  
    84  					break
    85  				}
    86  			}
    87  
    88  			if !shouldKeep {
    89  				delete(tr.allTSIDs, ts)
    90  			}
    91  		}
    92  	}
    93  
    94  	return nil
    95  }
    96  
    97  // For all incoming tsids, always add tsid and groupid to stored tsids
    98  func (tr *AllMatchedTSIDs) BulkAddStar(rawTagValueToTSIDs map[string]map[uint64]struct{}) error {
    99  	var err error
   100  	for tagValue, tsids := range rawTagValueToTSIDs {
   101  		for id := range tsids {
   102  			buf, ok := tr.allTSIDs[id]
   103  			if !ok {
   104  				buf = bytebufferpool.Get()
   105  				tr.allTSIDs[id] = buf
   106  			}
   107  			_, err = buf.WriteString(tagValue)
   108  			if err != nil {
   109  				return err
   110  			}
   111  			_, err = buf.Write(TAG_VALUE_DELIMITER_BYTE)
   112  			if err != nil {
   113  				return err
   114  			}
   115  		}
   116  	}
   117  
   118  	return nil
   119  }
   120  
   121  // needs to know if first update or not
   122  func (tr *AllMatchedTSIDs) FinishBlock() error {
   123  	tr.first = false
   124  	return nil
   125  }
   126  
   127  /*
   128  TSO access is highly optimized for TSIDs to be accessed in increasing order
   129  
   130  This reorders the matched TSIDS in increasing order
   131  */
   132  func (tr *AllMatchedTSIDs) FinishAllMatches() {
   133  	keys := make([]uint64, 0, len(tr.allTSIDs))
   134  	for k := range tr.allTSIDs {
   135  		keys = append(keys, k)
   136  	}
   137  	sort.Slice(keys, func(i, j int) bool {
   138  		return keys[i] < keys[j]
   139  	})
   140  
   141  	retVal := make(map[uint64]*bytebufferpool.ByteBuffer, len(tr.allTSIDs))
   142  	for _, k := range keys {
   143  		retVal[k] = tr.allTSIDs[k]
   144  	}
   145  	tr.allTSIDs = retVal
   146  }
   147  
   148  func (tr *AllMatchedTSIDs) GetNumMatchedTSIDs() int {
   149  	return len(tr.allTSIDs)
   150  }
   151  
   152  // returns a map of tsid to groupid
   153  func (tr *AllMatchedTSIDs) GetAllTSIDs() map[uint64]*bytebufferpool.ByteBuffer {
   154  	return tr.allTSIDs
   155  }