github.com/google/fleetspeak@v0.1.15-0.20240426164851-4f31f62c1aea/fleetspeak/src/client/internal/monitoring/resource_usage_monitor_test.go (about)

     1  // Copyright 2017 Google Inc.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     https://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package monitoring
    16  
    17  import (
    18  	"context"
    19  	"fmt"
    20  	"sync"
    21  	"testing"
    22  	"time"
    23  
    24  	"google.golang.org/protobuf/proto"
    25  
    26  	"github.com/google/fleetspeak/fleetspeak/src/client/clitesting"
    27  
    28  	fspb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak"
    29  	mpb "github.com/google/fleetspeak/fleetspeak/src/common/proto/fleetspeak_monitoring"
    30  	tspb "google.golang.org/protobuf/types/known/timestamppb"
    31  )
    32  
    33  type fakeResourceUsageFetcher struct {
    34  	resourceUsageFetcherI
    35  	ruList           []ResourceUsage // Fake resource-usage data to cycle through
    36  	statusList       []string        // Fake debug status strings to cycle through
    37  	ruCTR, statusCTR int
    38  }
    39  
    40  func (f *fakeResourceUsageFetcher) ResourceUsageForPID(pid int) (*ResourceUsage, error) {
    41  	ret := f.ruList[f.ruCTR]
    42  	f.ruCTR++
    43  	f.ruCTR %= len(f.ruList)
    44  	return &ret, nil
    45  }
    46  
    47  func (f *fakeResourceUsageFetcher) DebugStatusForPID(pid int) (string, error) {
    48  	ret := f.statusList[f.statusCTR]
    49  	f.statusCTR++
    50  	f.statusCTR %= len(f.statusList)
    51  	return ret, nil
    52  }
    53  
    54  func (f *fakeResourceUsageFetcher) setResourceUsageData(ruList []ResourceUsage) {
    55  	f.ruList = ruList
    56  }
    57  
    58  func (f *fakeResourceUsageFetcher) setDebugStatusStrings(statusList []string) {
    59  	f.statusList = statusList
    60  }
    61  
    62  func TestResourceUsageMonitor(t *testing.T) {
    63  	start := time.Now()
    64  	sc := clitesting.MockServiceContext{
    65  		OutChan: make(chan *fspb.Message),
    66  	}
    67  	pid := 1234
    68  	fakeStart := int64(1234567890)
    69  	samplePeriod := 100 * time.Millisecond
    70  	sampleSize := 2
    71  	errChan := make(chan error)
    72  
    73  	fakeRU0 := ResourceUsage{
    74  		Timestamp:       start,
    75  		UserCPUMillis:   13.0,
    76  		SystemCPUMillis: 5.0,
    77  		ResidentMemory:  10000,
    78  		NumFDs:          8,
    79  	}
    80  	fakeRU1 := ResourceUsage{
    81  		Timestamp:       start.Add(10 * time.Millisecond),
    82  		UserCPUMillis:   27.0,
    83  		SystemCPUMillis: 9.0,
    84  		ResidentMemory:  30000,
    85  		NumFDs:          42,
    86  	}
    87  	ruf := fakeResourceUsageFetcher{}
    88  	ruf.setResourceUsageData([]ResourceUsage{fakeRU0, fakeRU1})
    89  	ruf.setDebugStatusStrings([]string{"Fake Debug Status 0", "Fake Debug Status 1"})
    90  
    91  	rum, err := New(&sc, ResourceUsageMonitorParams{
    92  		Scope:            "test-scope",
    93  		Pid:              pid,
    94  		ProcessStartTime: time.Unix(fakeStart, 0),
    95  		MaxSamplePeriod:  samplePeriod,
    96  		SampleSize:       sampleSize,
    97  		Err:              errChan,
    98  		ruf:              &ruf,
    99  	})
   100  	if err != nil {
   101  		t.Fatal(err)
   102  	}
   103  
   104  	var wg sync.WaitGroup
   105  	defer wg.Wait()
   106  
   107  	ctx, cancel := context.WithCancel(context.Background())
   108  	defer cancel()
   109  
   110  	wg.Add(1)
   111  	go func() {
   112  		wg.Done()
   113  		rum.Run(ctx)
   114  	}()
   115  
   116  	// Wait up to one minute for each resource-usage report.
   117  	timeout := 1 * time.Minute
   118  	timeoutWhen := time.After(timeout)
   119  
   120  	for protosReceived := 0; protosReceived < 2; protosReceived++ {
   121  		select {
   122  		case m := <-sc.OutChan:
   123  			got := &mpb.ResourceUsageData{}
   124  			if err := m.Data.UnmarshalTo(got); err != nil {
   125  				t.Fatalf("Unable to unmarshal ResourceUsageData: %v", err)
   126  			}
   127  			want := &mpb.ResourceUsageData{
   128  				Scope:            "test-scope",
   129  				Pid:              int64(pid),
   130  				ProcessStartTime: &tspb.Timestamp{Seconds: fakeStart},
   131  				ResourceUsage: &mpb.AggregatedResourceUsage{
   132  					MeanUserCpuRate:    1400.0,
   133  					MaxUserCpuRate:     1400.0,
   134  					MeanSystemCpuRate:  400.0,
   135  					MaxSystemCpuRate:   400.0,
   136  					MeanResidentMemory: 20000.0,
   137  					MaxResidentMemory:  30000,
   138  					MeanNumFds:         25,
   139  					MaxNumFds:          42,
   140  				},
   141  				DebugStatus:   fmt.Sprintf("Fake Debug Status %d", protosReceived),
   142  				DataTimestamp: got.DataTimestamp,
   143  			}
   144  			if !proto.Equal(got, want) {
   145  				t.Errorf(
   146  					"Resource-usage proto %d received is not what we expect; got:\n%v\nwant:\n%v", protosReceived, got.String(), want.String())
   147  			}
   148  
   149  			if err := got.DataTimestamp.CheckValid(); err != nil {
   150  				t.Fatal(err)
   151  			}
   152  			timestamp := got.DataTimestamp.AsTime()
   153  			if timestamp.Before(start) {
   154  				t.Errorf(
   155  					"Timestamp for resource-usage proto %d [%v] is expected to be after %v",
   156  					protosReceived, timestamp, start)
   157  			}
   158  		case err := <-errChan:
   159  			t.Error(err)
   160  		case <-timeoutWhen:
   161  			t.Fatalf("Received %d resource-usage protos after %v. Wanted 2.", protosReceived, timeout)
   162  		}
   163  	}
   164  }