signals.go

Documentation: k8s.io/kubernetes/test/utils/ktesting

     1  /*
     2  Copyright 2023 The Kubernetes Authors.
     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 ktesting
    18  
    19  import (
    20  	"context"
    21  	"errors"
    22  	"io"
    23  	"os"
    24  	"os/signal"
    25  	"strings"
    26  	"sync"
    27  )
    28  
    29  var (
    30  	interruptCtx context.Context
    31  
    32  	defaultProgressReporter = new(progressReporter)
    33  	defaultSignalChannel    chan os.Signal
    34  )
    35  
    36  const ginkgoSpecContextKey = "GINKGO_SPEC_CONTEXT"
    37  
    38  type ginkgoReporter interface {
    39  	AttachProgressReporter(reporter func() string) func()
    40  }
    41  
    42  func init() {
    43  	// Setting up signals is intentionally done in an init function because
    44  	// then importing ktesting in a unit or integration test is sufficient
    45  	// to activate the signal behavior.
    46  	signalCtx, _ := signal.NotifyContext(context.Background(), os.Interrupt)
    47  	cancelCtx, cancel := context.WithCancelCause(context.Background())
    48  	go func() {
    49  		<-signalCtx.Done()
    50  		cancel(errors.New("received interrupt signal"))
    51  	}()
    52  
    53  	// This reimplements the contract between Ginkgo and Gomega for progress reporting.
    54  	// When using Ginkgo contexts, Ginkgo will implement it. This here is for "go test".
    55  	//
    56  	// nolint:staticcheck // It complains about using a plain string. This can only be fixed
    57  	// by Ginkgo and Gomega formalizing this interface and define a type (somewhere...
    58  	// probably cannot be in either Ginkgo or Gomega).
    59  	interruptCtx = context.WithValue(cancelCtx, ginkgoSpecContextKey, defaultProgressReporter)
    60  
    61  	defaultSignalChannel = make(chan os.Signal, 1)
    62  	// progressSignals will be empty on Windows.
    63  	if len(progressSignals) > 0 {
    64  		signal.Notify(defaultSignalChannel, progressSignals...)
    65  	}
    66  
    67  	// os.Stderr gets redirected by "go test". "go test -v" has to be
    68  	// used to see the output while a test runs.
    69  	defaultProgressReporter.setOutput(os.Stderr)
    70  	go defaultProgressReporter.run(interruptCtx, defaultSignalChannel)
    71  }
    72  
    73  type progressReporter struct {
    74  	mutex           sync.Mutex
    75  	reporterCounter int64
    76  	reporters       map[int64]func() string
    77  	out             io.Writer
    78  }
    79  
    80  var _ ginkgoReporter = &progressReporter{}
    81  
    82  func (p *progressReporter) setOutput(out io.Writer) io.Writer {
    83  	p.mutex.Lock()
    84  	defer p.mutex.Unlock()
    85  	oldOut := p.out
    86  	p.out = out
    87  	return oldOut
    88  }
    89  
    90  // AttachProgressReporter implements Gomega's contextWithAttachProgressReporter.
    91  func (p *progressReporter) AttachProgressReporter(reporter func() string) func() {
    92  	p.mutex.Lock()
    93  	defer p.mutex.Unlock()
    94  
    95  	// TODO (?): identify the caller and record that for dumpProgress.
    96  	p.reporterCounter++
    97  	id := p.reporterCounter
    98  	if p.reporters == nil {
    99  		p.reporters = make(map[int64]func() string)
   100  	}
   101  	p.reporters[id] = reporter
   102  	return func() {
   103  		p.detachProgressReporter(id)
   104  	}
   105  }
   106  
   107  func (p *progressReporter) detachProgressReporter(id int64) {
   108  	p.mutex.Lock()
   109  	defer p.mutex.Unlock()
   110  
   111  	delete(p.reporters, id)
   112  }
   113  
   114  func (p *progressReporter) run(ctx context.Context, progressSignalChannel chan os.Signal) {
   115  	for {
   116  		select {
   117  		case <-ctx.Done():
   118  			return
   119  		case <-progressSignalChannel:
   120  			p.dumpProgress()
   121  		}
   122  	}
   123  }
   124  
   125  // dumpProgress is less useful than the Ginkgo progress report. We can't fix
   126  // that we don't know which tests are currently running and instead have to
   127  // rely on "go test -v" for that.
   128  //
   129  // But perhaps dumping goroutines and their callstacks is useful anyway?  TODO:
   130  // look at how Ginkgo does it and replicate some of it.
   131  func (p *progressReporter) dumpProgress() {
   132  	p.mutex.Lock()
   133  	defer p.mutex.Unlock()
   134  
   135  	var buffer strings.Builder
   136  	buffer.WriteString("You requested a progress report.\n")
   137  	if len(p.reporters) == 0 {
   138  		buffer.WriteString("Currently there is no information about test progress available.\n")
   139  	}
   140  	for _, reporter := range p.reporters {
   141  		report := reporter()
   142  		buffer.WriteRune('\n')
   143  		buffer.WriteString(report)
   144  		if !strings.HasSuffix(report, "\n") {
   145  			buffer.WriteRune('\n')
   146  		}
   147  	}
   148  
   149  	_, _ = p.out.Write([]byte(buffer.String()))
   150  }
   151
View as plain text