gvisor.dev/gvisor@v0.0.0-20240520182842-f9d4d51c7e0f/pkg/tcpip/link/qdisc/fifo/qdisc_test.go (about)

     1  // Copyright 2021 The gVisor Authors.
     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  //     http://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 qdisc_test
    16  
    17  import (
    18  	"math/rand"
    19  	"os"
    20  	"testing"
    21  	"time"
    22  
    23  	"gvisor.dev/gvisor/pkg/buffer"
    24  	"gvisor.dev/gvisor/pkg/refs"
    25  	"gvisor.dev/gvisor/pkg/sync"
    26  	"gvisor.dev/gvisor/pkg/tcpip"
    27  	"gvisor.dev/gvisor/pkg/tcpip/link/qdisc/fifo"
    28  	"gvisor.dev/gvisor/pkg/tcpip/stack"
    29  )
    30  
    31  var _ stack.LinkWriter = (*countWriter)(nil)
    32  
    33  // countWriter implements LinkWriter.
    34  type countWriter struct {
    35  	mu             sync.Mutex
    36  	packetsWritten int
    37  	packetsWanted  int
    38  	done           chan struct{}
    39  }
    40  
    41  func (cw *countWriter) WritePackets(pkts stack.PacketBufferList) (int, tcpip.Error) {
    42  	cw.mu.Lock()
    43  	defer cw.mu.Unlock()
    44  	cw.packetsWritten += pkts.Len()
    45  	// Opt out of using the done channel if packetsWanted is not set.
    46  	if cw.packetsWanted > 0 && cw.packetsWritten == cw.packetsWanted {
    47  		close(cw.done)
    48  	}
    49  	return pkts.Len(), nil
    50  }
    51  
    52  // In b/209690936, fast simultaneous writes on qdisc will cause panics. This test
    53  // reproduces the behavior shown in that bug.
    54  func TestFastSimultaneousWrites(t *testing.T) {
    55  	lower := &countWriter{}
    56  	linkEP := fifo.New(lower, 16, 1000)
    57  
    58  	v := make([]byte, 1)
    59  
    60  	// Simulate many simultaneous writes from various goroutines, similar to TCP's sendTCPBatch().
    61  	nWriters := 100
    62  	nWrites := 100
    63  	var wg sync.WaitGroup
    64  	for i := 0; i < nWriters; i++ {
    65  		wg.Add(1)
    66  		go func() {
    67  			defer wg.Done()
    68  			for j := 0; j < nWrites; j++ {
    69  				pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
    70  					Payload: buffer.MakeWithData(v),
    71  				})
    72  				pkt.Hash = rand.Uint32()
    73  				linkEP.WritePacket(pkt)
    74  				pkt.DecRef()
    75  			}
    76  		}()
    77  	}
    78  	wg.Wait()
    79  	linkEP.Close()
    80  }
    81  
    82  func TestWriteRefusedAfterClosed(t *testing.T) {
    83  	linkEp := fifo.New(nil, 1, 2)
    84  
    85  	linkEp.Close()
    86  	err := linkEp.WritePacket(nil)
    87  	_, ok := err.(*tcpip.ErrClosedForSend)
    88  	if !ok {
    89  		t.Errorf("got err = %s, want %s", err, &tcpip.ErrClosedForSend{})
    90  	}
    91  }
    92  
    93  func TestWriteMorePacketsThanBatchSize(t *testing.T) {
    94  	tc := []int{fifo.BatchSize + 1, fifo.BatchSize*2 + 1}
    95  	v := make([]byte, 1)
    96  
    97  	for _, want := range tc {
    98  		done := make(chan struct{})
    99  		lower := &countWriter{done: done, packetsWanted: want}
   100  		linkEp := fifo.New(lower, 1, 1000)
   101  		for i := 0; i < want; i++ {
   102  			pkt := stack.NewPacketBuffer(stack.PacketBufferOptions{
   103  				Payload: buffer.MakeWithData(v),
   104  			})
   105  			linkEp.WritePacket(pkt)
   106  			pkt.DecRef()
   107  		}
   108  		select {
   109  		case <-done:
   110  		case <-time.After(1 * time.Second):
   111  			t.Fatalf("expected %d packets, but got only %d", want, lower.packetsWritten)
   112  		}
   113  		linkEp.Close()
   114  	}
   115  }
   116  
   117  func TestMain(m *testing.M) {
   118  	refs.SetLeakMode(refs.LeaksPanic)
   119  	code := m.Run()
   120  	refs.DoLeakCheck()
   121  	os.Exit(code)
   122  }