<s> """<STR_LIT>""" <EOL> import bisect <EOL> import functools <EOL> import threading <EOL> import time <EOL> from grr . lib import utils <EOL> from grr . lib . rdfvalues import structs <EOL> from grr . proto import jobs_pb2 <EOL> class Timed ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , varname , fields = None ) : <EOL> self . varname = varname <EOL> self . fields = fields or [ ] <EOL> def __call__ ( self , func ) : <EOL> @ functools . wraps ( func ) <EOL> def Decorated ( * args , ** kwargs ) : <EOL> start_time = time . time ( ) <EOL> res = None <EOL> try : <EOL> res = func ( * args , ** kwargs ) <EOL> finally : <EOL> total_time = time . time ( ) - start_time <EOL> STATS . RecordEvent ( self . varname , total_time , fields = self . fields ) <EOL> return res <EOL> return Decorated <EOL> class Counted ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , varname , fields = None ) : <EOL> self . varname = varname <EOL> self . fields = fields or [ ] <EOL> def __call__ ( self , func ) : <EOL> @ functools . wraps ( func ) <EOL> def Decorated ( * args , ** kwargs ) : <EOL> try : <EOL> res = func ( * args , ** kwargs ) <EOL> finally : <EOL> STATS . IncrementCounter ( self . varname , fields = self . fields ) <EOL> return res <EOL> return Decorated <EOL> class CountingExceptionMixin ( object ) : <EOL> """<STR_LIT>""" <EOL> counter = None <EOL> fields = [ ] <EOL> def __init__ ( self , * args , ** kwargs ) : <EOL> if self . counter : <EOL> STATS . IncrementCounter ( self . counter , fields = self . fields ) <EOL> super ( CountingExceptionMixin , self ) . __init__ ( * args , ** kwargs ) <EOL> class _Metric ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self , fields_defs , docstring , units ) : <EOL> self . fields_defs = fields_defs <EOL> self . docstring = docstring <EOL> self . units = units <EOL> self . _values = { } <EOL> def _FieldsToKey ( self , fields ) : <EOL> if not fields : <EOL> return ( "<STR_LIT>" , ) <EOL> else : <EOL> return tuple ( fields ) <EOL> def _DefaultValue ( self ) : <EOL> return "<STR_LIT>" <EOL> def Get ( self , fields = None ) : <EOL> """<STR_LIT>""" <EOL> if self . fields_defs is None and fields is not None : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" % fields ) <EOL> if self . fields_defs is not None and fields is None : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" % self . fields_defs ) <EOL> if ( self . fields_defs is not None and fields is not None and <EOL> len ( self . fields_defs ) != len ( fields ) ) : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" % ( len ( self . fields_defs ) , <EOL> self . fields_defs , <EOL> len ( fields ) , <EOL> fields ) ) <EOL> try : <EOL> return self . _values [ self . _FieldsToKey ( fields ) ] <EOL> except KeyError : <EOL> return self . _DefaultValue ( ) <EOL> def ListFieldsValues ( self ) : <EOL> """<STR_LIT>""" <EOL> if self . fields_defs : <EOL> return self . _values . iterkeys ( ) <EOL> else : <EOL> return [ ] <EOL> class _CounterMetric ( _Metric ) : <EOL> """<STR_LIT>""" <EOL> def _DefaultValue ( self ) : <EOL> return <NUM_LIT:0> <EOL> def Increment ( self , delta , fields = None ) : <EOL> """<STR_LIT>""" <EOL> if delta < <NUM_LIT:0> : <EOL> raise ValueError ( "<STR_LIT>" % delta ) <EOL> key = self . _FieldsToKey ( fields ) <EOL> if key in self . _values : <EOL> self . _values [ key ] += delta <EOL> else : <EOL> self . _values [ key ] = delta <EOL> class Distribution ( structs . RDFProtoStruct ) : <EOL> """<STR_LIT>""" <EOL> protobuf = jobs_pb2 . Distribution <EOL> def __init__ ( self , initializer = None , age = None , bins = None ) : <EOL> if initializer and bins : <EOL> raise ValueError ( "<STR_LIT>" <EOL> "<STR_LIT>" ) <EOL> super ( Distribution , self ) . __init__ ( initializer = initializer , age = age ) <EOL> if bins : <EOL> self . bins = [ - float ( "<STR_LIT>" ) ] + bins <EOL> self . heights = [ <NUM_LIT:0> ] * len ( self . bins ) <EOL> def Record ( self , value ) : <EOL> """<STR_LIT>""" <EOL> self . sum += value <EOL> self . count += <NUM_LIT:1> <EOL> pos = bisect . bisect ( self . bins , value ) - <NUM_LIT:1> <EOL> if pos < <NUM_LIT:0> : <EOL> pos = <NUM_LIT:0> <EOL> elif pos == len ( self . bins ) : <EOL> pos = len ( self . bins ) - <NUM_LIT:1> <EOL> self . heights [ pos ] += <NUM_LIT:1> <EOL> @ property <EOL> def bins_heights ( self ) : <EOL> return dict ( zip ( self . bins , self . heights ) ) <EOL> class _EventMetric ( _Metric ) : <EOL> """<STR_LIT>""" <EOL> def _DefaultValue ( self ) : <EOL> return Distribution ( bins = self . _bins ) <EOL> def __init__ ( self , bins , fields , docstring , units ) : <EOL> self . _bins = bins or [ <NUM_LIT:0.0> , <NUM_LIT:0.1> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT> , <NUM_LIT:0.5> , <NUM_LIT> , <NUM_LIT:1> , <EOL> <NUM_LIT> , <NUM_LIT:2> , <NUM_LIT> , <NUM_LIT:3> , <NUM_LIT:4> , <NUM_LIT:5> , <NUM_LIT:6> , <NUM_LIT:7> , <NUM_LIT:8> , <NUM_LIT:9> , <NUM_LIT:10> , <EOL> <NUM_LIT:15> , <NUM_LIT:20> , <NUM_LIT:50> , <NUM_LIT:100> ] <EOL> super ( _EventMetric , self ) . __init__ ( fields , docstring , units ) <EOL> def Record ( self , value , fields = None ) : <EOL> """<STR_LIT>""" <EOL> key = self . _FieldsToKey ( fields ) <EOL> try : <EOL> entry = self . _values [ key ] <EOL> except KeyError : <EOL> entry = Distribution ( bins = self . _bins ) <EOL> self . _values [ key ] = entry <EOL> entry . Record ( value ) <EOL> class _GaugeMetric ( _Metric ) : <EOL> """<STR_LIT>""" <EOL> def _DefaultValue ( self ) : <EOL> return self . value_type ( ) <EOL> def __init__ ( self , value_type , fields , docstring , units ) : <EOL> self . value_type = value_type <EOL> super ( _GaugeMetric , self ) . __init__ ( fields , docstring , units ) <EOL> def Set ( self , value , fields = None ) : <EOL> """<STR_LIT>""" <EOL> key = self . _FieldsToKey ( fields ) <EOL> self . _values [ key ] = self . value_type ( value ) <EOL> def SetCallback ( self , callback , fields = None ) : <EOL> """<STR_LIT>""" <EOL> key = self . _FieldsToKey ( fields ) <EOL> self . _values [ key ] = callback <EOL> def Get ( self , fields = None ) : <EOL> """<STR_LIT>""" <EOL> result = super ( _GaugeMetric , self ) . Get ( fields = fields ) <EOL> if hasattr ( result , "<STR_LIT>" ) : <EOL> return result ( ) <EOL> else : <EOL> return result <EOL> class MetricFieldDefinition ( structs . RDFProtoStruct ) : <EOL> """<STR_LIT>""" <EOL> protobuf = jobs_pb2 . MetricFieldDefinition <EOL> class MetricMetadata ( structs . RDFProtoStruct ) : <EOL> """<STR_LIT>""" <EOL> protobuf = jobs_pb2 . MetricMetadata <EOL> def DefaultValue ( self ) : <EOL> if self . value_type == self . ValueType . INT : <EOL> return <NUM_LIT:0> <EOL> elif self . value_type == self . ValueType . FLOAT : <EOL> return <NUM_LIT:0.0> <EOL> elif self . value_type == self . ValueType . STR : <EOL> return "<STR_LIT>" <EOL> else : <EOL> return Distribution ( ) <EOL> MetricType = MetricMetadata . MetricType <EOL> MetricUnits = MetricMetadata . MetricUnits <EOL> class StatsCollector ( object ) : <EOL> """<STR_LIT>""" <EOL> def __init__ ( self ) : <EOL> self . _metrics = { } <EOL> self . lock = threading . RLock ( ) <EOL> self . _metrics_metadata = { } <EOL> @ staticmethod <EOL> def ValueTypeToMetricValueType ( value_type ) : <EOL> """<STR_LIT>""" <EOL> if value_type in ( int , long ) : <EOL> value_type = MetricMetadata . ValueType . INT <EOL> elif value_type == str : <EOL> value_type = MetricMetadata . ValueType . STR <EOL> elif value_type == float : <EOL> value_type = MetricMetadata . ValueType . FLOAT <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" % value_type ) <EOL> return value_type <EOL> @ staticmethod <EOL> def FieldsToFieldsDefinitions ( fields ) : <EOL> """<STR_LIT>""" <EOL> if not fields : <EOL> return [ ] <EOL> result = [ ] <EOL> for field_name , field_type in fields : <EOL> if field_type in ( int , long ) : <EOL> field_type = MetricFieldDefinition . FieldType . INT <EOL> elif field_type == str : <EOL> field_type = MetricFieldDefinition . FieldType . STR <EOL> else : <EOL> raise ValueError ( "<STR_LIT>" % field_type ) <EOL> result . append ( MetricFieldDefinition ( field_name = field_name , <EOL> field_type = field_type ) ) <EOL> return result <EOL> @ utils . Synchronized <EOL> def RegisterCounterMetric ( self , varname , fields = None , docstring = None , <EOL> units = None ) : <EOL> """<STR_LIT>""" <EOL> self . _metrics [ varname ] = _CounterMetric ( fields , docstring , units ) <EOL> self . _metrics_metadata [ varname ] = MetricMetadata ( <EOL> varname = varname , metric_type = MetricMetadata . MetricType . COUNTER , <EOL> value_type = MetricMetadata . ValueType . INT , <EOL> fields_defs = self . FieldsToFieldsDefinitions ( fields ) , <EOL> docstring = docstring , units = units ) <EOL> @ utils . Synchronized <EOL> def IncrementCounter ( self , varname , delta = <NUM_LIT:1> , fields = None ) : <EOL> """<STR_LIT>""" <EOL> self . _metrics [ varname ] . Increment ( delta , fields ) <EOL> @ utils . Synchronized <EOL> def RegisterEventMetric ( self , varname , bins = None , fields = None , <EOL> docstring = None , units = None ) : <EOL> """<STR_LIT>""" <EOL> self . _metrics [ varname ] = _EventMetric ( bins , fields , docstring , units ) <EOL> self . _metrics_metadata [ varname ] = MetricMetadata ( <EOL> varname = varname , metric_type = MetricMetadata . MetricType . EVENT , <EOL> value_type = MetricMetadata . ValueType . DISTRIBUTION , <EOL> fields_defs = self . FieldsToFieldsDefinitions ( fields ) , <EOL> docstring = docstring , units = units ) <EOL> @ utils . Synchronized <EOL> def RecordEvent ( self , varname , value , fields = None ) : <EOL> """<STR_LIT>""" <EOL> self . _metrics [ varname ] . Record ( value , fields ) <EOL> @ utils . Synchronized <EOL> def RegisterGaugeMetric ( self , varname , value_type , fields = None , <EOL> docstring = None , units = None ) : <EOL> """<STR_LIT>""" <EOL> self . _metrics [ varname ] = _GaugeMetric ( value_type , fields , docstring , units ) <EOL> self . _metrics_metadata [ varname ] = MetricMetadata ( <EOL> varname = varname , metric_type = MetricMetadata . MetricType . GAUGE , <EOL> value_type = self . ValueTypeToMetricValueType ( value_type ) , <EOL> fields_defs = self . FieldsToFieldsDefinitions ( fields ) , <EOL> docstring = docstring , units = units ) <EOL> @ utils . Synchronized <EOL> def SetGaugeValue ( self , varname , value , fields = None ) : <EOL> """<STR_LIT>""" <EOL> self . _metrics [ varname ] . Set ( value , fields ) <EOL> @ utils . Synchronized <EOL> def SetGaugeCallback ( self , varname , callback , fields = None ) : <EOL> """<STR_LIT>""" <EOL> self . _metrics [ varname ] . SetCallback ( callback , fields ) <EOL> def GetMetricMetadata ( self , varname ) : <EOL> """<STR_LIT>""" <EOL> return self . _metrics_metadata [ varname ] <EOL> def GetAllMetricsMetadata ( self ) : <EOL> """<STR_LIT>""" <EOL> return self . _metrics_metadata . copy ( ) <EOL> def GetMetricFields ( self , varname ) : <EOL> """<STR_LIT>""" <EOL> return self . _metrics [ varname ] . ListFieldsValues ( ) <EOL> def GetMetricValue ( self , varname , fields = None ) : <EOL> """<STR_LIT>""" <EOL> return self . _metrics [ varname ] . Get ( fields ) <EOL> STATS = None </s>
