github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/agent/session_test.go (about)

     1  package agent
     2  
     3  import (
     4  	"os"
     5  	"time"
     6  
     7  	. "github.com/onsi/ginkgo/v2"
     8  	. "github.com/onsi/gomega"
     9  	"github.com/pyroscope-io/pyroscope/pkg/agent/spy"
    10  	"github.com/pyroscope-io/pyroscope/pkg/agent/upstream"
    11  	"github.com/sirupsen/logrus"
    12  )
    13  
    14  type upstreamMock struct {
    15  	uploads []*upstream.UploadJob
    16  }
    17  
    18  func (*upstreamMock) Start() {}
    19  func (*upstreamMock) Stop()  {}
    20  
    21  func (u *upstreamMock) Upload(j *upstream.UploadJob) {
    22  	u.uploads = append(u.uploads, j)
    23  }
    24  
    25  var _ = Describe("agent.Session", func() {
    26  	Describe("NewSession", func() {
    27  		It("creates a new session and performs chunking", func() {
    28  			done := make(chan interface{})
    29  			go func() {
    30  				defer GinkgoRecover()
    31  				u := &upstreamMock{}
    32  				uploadRate := 200 * time.Millisecond
    33  				s, _ := NewSession(SessionConfig{
    34  					Upstream:         u,
    35  					AppName:          "test-app",
    36  					ProfilingTypes:   []spy.ProfileType{spy.ProfileCPU},
    37  					SpyName:          "debugspy",
    38  					SampleRate:       100,
    39  					UploadRate:       uploadRate,
    40  					Pid:              os.Getpid(),
    41  					WithSubprocesses: true,
    42  					Logger:           logrus.StandardLogger(),
    43  				})
    44  				now := time.Now()
    45  				time.Sleep(now.Truncate(uploadRate).Add(uploadRate + 10*time.Millisecond).Sub(now))
    46  				err := s.Start()
    47  
    48  				Expect(err).ToNot(HaveOccurred())
    49  				time.Sleep(500 * time.Millisecond)
    50  				s.Stop()
    51  
    52  				Expect(u.uploads).To(HaveLen(3))
    53  				u.uploads[0].Trie.Iterate(func(name []byte, val uint64) {
    54  					Expect(val).To(BeNumerically("~", 19, 2))
    55  				})
    56  				u.uploads[1].Trie.Iterate(func(name []byte, val uint64) {
    57  					Expect(val).To(BeNumerically("~", 20, 2))
    58  				})
    59  				u.uploads[2].Trie.Iterate(func(name []byte, val uint64) {
    60  					Expect(val).To(BeNumerically("~", 11, 2))
    61  				})
    62  				close(done)
    63  				Expect(u.uploads[0].Name).To(Equal("test-app.cpu{}"))
    64  			}()
    65  			Eventually(done, 3).Should(BeClosed())
    66  		})
    67  
    68  		When("tags specified", func() {
    69  			It("name ", func() {
    70  				u := &upstreamMock{}
    71  				uploadRate := 200 * time.Millisecond
    72  				c := SessionConfig{
    73  					Upstream:         u,
    74  					AppName:          "test-app{bar=xxx}",
    75  					ProfilingTypes:   []spy.ProfileType{spy.ProfileCPU},
    76  					SpyName:          "debugspy",
    77  					SampleRate:       100,
    78  					UploadRate:       uploadRate,
    79  					Pid:              os.Getpid(),
    80  					Logger:           logrus.StandardLogger(),
    81  					WithSubprocesses: true,
    82  					Tags: map[string]string{
    83  						"foo": "bar",
    84  						"baz": "qux",
    85  					},
    86  				}
    87  
    88  				s, _ := NewSession(c)
    89  				now := time.Now()
    90  				time.Sleep(now.Truncate(uploadRate).Add(uploadRate + 10*time.Millisecond).Sub(now))
    91  				err := s.Start()
    92  				Expect(err).ToNot(HaveOccurred())
    93  				time.Sleep(500 * time.Millisecond)
    94  				s.Stop()
    95  
    96  				Expect(u.uploads).To(HaveLen(3))
    97  				Expect(u.uploads[0].Name).To(Equal("test-app.cpu{bar=xxx,baz=qux,foo=bar}"))
    98  			})
    99  		})
   100  	})
   101  	Describe("Session", func() {
   102  		When("tags added", func() {
   103  			It("name ", func() {
   104  				u := &upstreamMock{}
   105  				uploadRate := 200 * time.Millisecond
   106  				c := SessionConfig{
   107  					Upstream:         u,
   108  					AppName:          "test-app",
   109  					ProfilingTypes:   []spy.ProfileType{spy.ProfileCPU},
   110  					SpyName:          "debugspy",
   111  					SampleRate:       100,
   112  					UploadRate:       uploadRate,
   113  					Pid:              os.Getpid(),
   114  					WithSubprocesses: true,
   115  					Logger:           logrus.StandardLogger(),
   116  					Tags: map[string]string{
   117  						"foo": "bar",
   118  						"baz": "qux",
   119  					},
   120  				}
   121  
   122  				s, _ := NewSession(c)
   123  				now := time.Now()
   124  				time.Sleep(now.Truncate(uploadRate).Add(uploadRate + 10*time.Millisecond).Sub(now))
   125  				err := s.Start()
   126  				Expect(err).ToNot(HaveOccurred())
   127  				s.SetTags(map[string]string{"bar": "xxx"})
   128  				time.Sleep(500 * time.Millisecond)
   129  				s.Stop()
   130  
   131  				Expect(u.uploads).To(HaveLen(3))
   132  				Expect(u.uploads[0].Name).To(Equal("test-app.cpu{bar=xxx,baz=qux,foo=bar}"))
   133  			})
   134  		})
   135  		When("tags removed", func() {
   136  			It("name ", func() {
   137  				u := &upstreamMock{}
   138  				uploadRate := 200 * time.Millisecond
   139  				c := SessionConfig{
   140  					Upstream:         u,
   141  					AppName:          "test-app{bar=xxx}",
   142  					ProfilingTypes:   []spy.ProfileType{spy.ProfileCPU},
   143  					SpyName:          "debugspy",
   144  					SampleRate:       100,
   145  					UploadRate:       uploadRate,
   146  					Pid:              os.Getpid(),
   147  					WithSubprocesses: true,
   148  					Logger:           logrus.StandardLogger(),
   149  					Tags: map[string]string{
   150  						"foo": "bar",
   151  					},
   152  				}
   153  
   154  				s, _ := NewSession(c)
   155  				now := time.Now()
   156  				time.Sleep(now.Truncate(uploadRate).Add(uploadRate + 10*time.Millisecond).Sub(now))
   157  				err := s.Start()
   158  				Expect(err).ToNot(HaveOccurred())
   159  				err = s.RemoveTags("bar")
   160  				Expect(err).ToNot(HaveOccurred())
   161  				err = s.RemoveTags("foo")
   162  				Expect(err).ToNot(HaveOccurred())
   163  				time.Sleep(500 * time.Millisecond)
   164  				s.Stop()
   165  
   166  				Expect(u.uploads).To(HaveLen(3))
   167  				Expect(u.uploads[0].Name).To(Equal("test-app.cpu{}"))
   168  			})
   169  		})
   170  	})
   171  })