github.com/anth0d/nomad@v0.0.0-20221214183521-ae3a0a2cad06/client/allocrunner/taskrunner/logmon_hook_test.go (about)

     1  package taskrunner
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"net"
     7  	"testing"
     8  
     9  	plugin "github.com/hashicorp/go-plugin"
    10  	"github.com/hashicorp/nomad/ci"
    11  	"github.com/hashicorp/nomad/client/allocrunner/interfaces"
    12  	"github.com/hashicorp/nomad/helper/testlog"
    13  	"github.com/hashicorp/nomad/nomad/mock"
    14  	pstructs "github.com/hashicorp/nomad/plugins/shared/structs"
    15  	"github.com/stretchr/testify/require"
    16  	"golang.org/x/exp/maps"
    17  )
    18  
    19  // Statically assert the logmon hook implements the expected interfaces
    20  var _ interfaces.TaskPrestartHook = (*logmonHook)(nil)
    21  var _ interfaces.TaskStopHook = (*logmonHook)(nil)
    22  
    23  // TestTaskRunner_LogmonHook_LoadReattach unit tests loading logmon reattach
    24  // config from persisted hook state.
    25  func TestTaskRunner_LogmonHook_LoadReattach(t *testing.T) {
    26  	ci.Parallel(t)
    27  
    28  	// No hook data should return nothing
    29  	cfg, err := reattachConfigFromHookData(nil)
    30  	require.Nil(t, cfg)
    31  	require.NoError(t, err)
    32  
    33  	// Hook data without the appropriate key should return nothing
    34  	cfg, err = reattachConfigFromHookData(map[string]string{"foo": "bar"})
    35  	require.Nil(t, cfg)
    36  	require.NoError(t, err)
    37  
    38  	// Create a realistic reattach config and roundtrip it
    39  	addr, err := net.ResolveTCPAddr("tcp", "127.0.0.1:0")
    40  	require.NoError(t, err)
    41  
    42  	orig := &plugin.ReattachConfig{
    43  		Protocol: plugin.ProtocolGRPC,
    44  		Addr:     addr,
    45  		Pid:      4,
    46  	}
    47  	origJSON, err := json.Marshal(pstructs.ReattachConfigFromGoPlugin(orig))
    48  	require.NoError(t, err)
    49  
    50  	cfg, err = reattachConfigFromHookData(map[string]string{
    51  		logmonReattachKey: string(origJSON),
    52  	})
    53  	require.NoError(t, err)
    54  
    55  	require.Equal(t, orig, cfg)
    56  }
    57  
    58  // TestTaskRunner_LogmonHook_StartStop asserts that a new logmon is created the
    59  // first time Prestart is called, reattached to on subsequent restarts, and
    60  // killed on Stop.
    61  func TestTaskRunner_LogmonHook_StartStop(t *testing.T) {
    62  	ci.Parallel(t)
    63  
    64  	alloc := mock.BatchAlloc()
    65  	task := alloc.Job.TaskGroups[0].Tasks[0]
    66  
    67  	dir := t.TempDir()
    68  
    69  	hookConf := newLogMonHookConfig(task.Name, dir)
    70  	runner := &TaskRunner{logmonHookConfig: hookConf}
    71  	hook := newLogMonHook(runner, testlog.HCLogger(t))
    72  
    73  	req := interfaces.TaskPrestartRequest{
    74  		Task: task,
    75  	}
    76  	resp := interfaces.TaskPrestartResponse{}
    77  
    78  	// First prestart should set reattach key but never be Done as it needs
    79  	// to rerun on agent restarts to reattach.
    80  	require.NoError(t, hook.Prestart(context.Background(), &req, &resp))
    81  	defer hook.Stop(context.Background(), nil, nil)
    82  
    83  	require.False(t, resp.Done)
    84  	origHookData := resp.State[logmonReattachKey]
    85  	require.NotEmpty(t, origHookData)
    86  
    87  	// Running prestart again should effectively noop as it reattaches to
    88  	// the running logmon.
    89  	req.PreviousState = map[string]string{
    90  		logmonReattachKey: origHookData,
    91  	}
    92  	require.NoError(t, hook.Prestart(context.Background(), &req, &resp))
    93  	require.False(t, resp.Done)
    94  	origHookData = resp.State[logmonReattachKey]
    95  	require.Equal(t, origHookData, req.PreviousState[logmonReattachKey])
    96  
    97  	// Running stop should shutdown logmon
    98  	stopReq := interfaces.TaskStopRequest{
    99  		ExistingState: maps.Clone(resp.State),
   100  	}
   101  	require.NoError(t, hook.Stop(context.Background(), &stopReq, nil))
   102  }