o
    i&                     @   s:  d dl mZ d dlmZ d dlmZ d dlmZ d dlmZ d dl	m
Z
 d dlmZ d dlmZ d d	lmZ d d
lmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ d dlmZ eeZdZ e!e ej" e#dZ$G dd de%eZ&eG dd dZ'G dd de(Z)G dd dej*Z+dS )    )	dataclass)Enum)Any)Optional)quote)	di_config)LogSignalJsonEncoder)SignalQueue)SnapshotJsonEncoder)metrics)SignalCollector)SignalTrack)agent)logger)
get_logger)	connector)fibonacci_backoff_with_jitterunsupported_agentuploaderc                   @   s   e Zd ZdZdZdZdZdS )UploaderProductzUploader products.dynamic_instrumentationexception_replayzcode_origin.span.entryN)__name__
__module____qualname____doc__DEBUGGEREXCEPTION_REPLAYCODE_ORIGIN_SPAN_ENTRY r   r   O/home/ubuntu/.local/lib/python3.10/site-packages/ddtrace/debugging/_uploader.pyr      s
    r   c                   @   s2   e Zd ZU eed< eed< eed< dZeed< dS )UploaderTracktrackendpointqueueTenabledN)	r   r   r   r   __annotations__strr	   r%   boolr   r   r   r    r!   %   s
   
 r!   c                   @   s   e Zd ZdZdS )SignalUploaderErrorzSignal uploader error.N)r   r   r   r   r   r   r   r    r)   -   s    r)   c                       s4  e Zd ZU dZdZed  ed< e Zee	 ed< e Z
ee ed< eZeZdZd(dee ddf fd	d
Zdee defddZdededdfddZdededdfddZd)ddZd)ddZd)ddZdeddfddZd)dd ZeZ e!dee fd!d"Z"e!d#e	ddfd$d%Z#e!d#e	ddfd&d'Z$  Z%S )*SignalUploaderzSignal uploader.

    This class implements an interface with the debugger signal intake for both
    the debugger and the events platform.
    N	_instance	_products_agent_endpoints   intervalreturnc              
      s,  t  |d ur	|ntj tjrtjrdttj nd | _}tj	t
tj	d| | jttj| jddtjt
tjd| | jttj| jddi| _| dd | j D | _dd	d
| _ttjtjd| _td| j d| j  d | jd| j| _t dtjdd | j D | j d| _!d S )Nz?ddtags= /debugger/v2/inputencoderon_full)r"   r#   r$   c                 S      i | ]\}}||j qS r   r$   .0tutr   r   r    
<dictcomp>X       z+SignalUploader.__init__.<locals>.<dictcomp>zapplication/json; charset=utf-8z
text/plain)zContent-typeAccept)timeoutg-?gS?   )initial_waitattemptszBSignal uploader initialized (url: %s, endpoints: %s, interval: %f)c                 S   r6   r   )r#   r8   r   r   r    r<   i   r=   F)"super__init__r   upload_interval_seconds_tags_in_qstagsr   _endpoint_suffixr   LOGSr!   	__queue__r   service_name_on_buffer_fullSNAPSHOTr
   _tracks__collector__items
_collector_headersr   _intake_urlupload_timeout_connectr   r/   RETRY_ATTEMPTS_write_write_with_backofflogdebug_flush_full)selfr/   endpoint_suffix	__class__r   r    rD   C   sD   
zSignalUploader.__init__
agent_infoc                 C   s   |d u rdS d|vrt d dS t|dg }| jtj }| jtj }d|_d|_d|v rDt d d| j	 |_
d| j	 |_
dS d|v r]t d d| j	 |_
d| j	 |_
dS d|_d|_t jtd	d
dd dS )NF	endpointsz=Unsupported Datadog agent detected. Please upgrade to 7.49.0.Tr2   z$Detected /debugger/v2/input endpoint/debugger/v1/diagnosticsz3Detected /debugger/v1/diagnostics endpoint fallbackdebuggerzUnsupported Datadog agent detected. Logs and snapshots from Dynamic Instrumentation/Exception Replay/Code Origin for Spans will not be uploaded. Please upgrade to version 7.49.0 or later)product	more_info)extra)rY   rZ   setgetrN   r   rI   rM   r%   rH   r#   warningUNSUPPORTED_AGENT)r\   r`   ra   
logs_tracksnapshot_trackr   r   r    
info_checko   s:   


zSignalUploader.info_checkpayloadr#   c                 C   s  zl|   ]}|jd||| jd | }d|j  krdk sCn td||j|  tj	ddt
|jid d	|jkrBd
}t|nt	d tdt| W d    W d S W d    W d S 1 sew   Y  W d S  tyt     ty   tjd|dd t	d Y d S w )NPOST)headers   i,  z0Failed to upload payload to endpoint %s: [%d] %rzupload.errorstatus)rG   i  zFailed to upload payloadzupload.successzupload.sizez&Failed to write payload to endpoint %sTexc_infoerror)rU   requestrR   getresponserr   rY   ru   readmeter	incrementr'   r)   distributionlen	Exception)r\   rn   r#   connrespmsgr   r   r    rW      s,   


&zSignalUploader._write_item_encodedc                 C   s   d| _ |   d S )NT)r[   upload)r\   r   r   r   r   r    rL      s   zSignalUploader._on_buffer_fullc                 C   s   |    dS )zUpload request.N)awake)r\   r   r   r    r      s   zSignalUploader.uploadc                 C   sB   | j  D ]}| j|jj| jd|_qdd | j  D | j_ dS )zReset the buffer on fork.r3   c                 S   r6   r   r7   r8   r   r   r    r<      r=   z(SignalUploader.reset.<locals>.<dictcomp>N)rN   valuesrJ   r$   _encoderrL   rP   rQ   )r\   r"   r   r   r    reset   s   zSignalUploader.resetc                 C   s8   d| j  }|| jtj _|| jtj _td| dS )z2Downgrade both tracks to the diagnostics endpoint.rb   z$Downgrading debugger endpoints to %sN)rH   rN   r   rI   r#   rM   rY   rZ   )r\   diagnostics_endpointr   r   r    _downgrade_to_diagnostics   s   z(SignalUploader._downgrade_to_diagnosticsr"   c                 C   s   |j   }d urS|jrU|\}}z| ||j td| W d S  tyB   |jdsA| 	  | ||j td| Y d S   t
yR   tjddd Y d S w d S d S )Nzbatch.cardinalityrb   zCannot upload payloadTrs   )r$   flushr%   rX   r#   ry   r{   r)   
startswithr   r}   rY   rZ   )r\   r"   datarn   countr   r   r    _flush_track   s   zSignalUploader._flush_trackc                 C   s~   | j rd| _ | j D ]}|j r| | q| j D ]}|jjr(| | q| jtj j	r7| jtj
 j	s=d}t|dS )z'Upload the buffer content to the agent.FzDebugger tracks not enabledN)r[   rN   r   r$   is_fullr   r   r   rM   r%   rI   
ValueError)r\   uploader_trackr"   r   r   r   r    online   s   


zSignalUploader.onlinec                 C   s   | j d ur	| j jS d S N)r+   rQ   )clsr   r   r    get_collector   s   zSignalUploader.get_collectorrd   c                 C   s>   || j v rd S | j | | jd u r|  | _| j  d S d S r   )r,   addr+   startr   rd   r   r   r    register   s   

zSignalUploader.registerc                 C   sF   || j vrd S | j | | j s| jd ur!| j  d | _d S d S d S r   )r,   remover+   stopr   r   r   r    
unregister   s   


zSignalUploader.unregisterr   )r0   N)&r   r   r   r   r+   r   r&   rg   r,   r   r-   r'   r	   rJ   r   rO   rV   floatrD   dictr(   rm   bytesrW   r   rL   r   r   r   r!   r   r   on_shutdownclassmethodr   r   r   __classcell__r   r   r^   r    r*   3   s0   
 ,*




r*   N),dataclassesr   enumr   typingr   r   urllib.parser   ddtrace.debugging._configr   ddtrace.debugging._encodingr   r	   r
   ddtrace.debugging._metricsr   #ddtrace.debugging._signal.collectorr   ddtrace.debugging._signal.modelr   ddtrace.internalr   r   ddtrace.internal.loggerr   ddtrace.internal.utils.httpr   ddtrace.internal.utils.retryr   r   rY   rj   set_tag_rate_limitHOUR	get_meterry   r'   r   r!   r}   r)   AgentCheckPeriodicServicer*   r   r   r   r    <module>   s4    
