Fix context issue

Signed-off-by: bpopovschi <[email protected]>

Author: Zyqsempai

runtime/service.go

@@ -506,8 +506,7 @@ func (s *service) StopVM(requestCtx context.Context, request *proto.StopVMReques
 	} else {
 		timeout = time.Duration(request.VMShutdownTimeoutSeconds) * time.Second
 	}
-	ctx, cancel := context.WithTimeout(requestCtx, timeout)
-	defer cancel()
+	timer := time.NewTimer(timeout)
 	defer logPanicAndDie(s.logger)
 	// If something goes wrong here, just shut down ungracefully. This eliminates some scenarios that would result
 	// in the user being unable to shut down the VM.
@@ -526,12 +525,12 @@ func (s *service) StopVM(requestCtx context.Context, request *proto.StopVMReques
 	shutdownCh := make(chan error)
 	go func() {
 		defer close(shutdownCh)
-		_, err = s.Shutdown(ctx, &taskAPI.ShutdownRequest{Now: true})
+		_, err = s.Shutdown(requestCtx, &taskAPI.ShutdownRequest{Now: true})
 		shutdownCh <- err
 	}()
 
 	select {
-	case <-ctx.Done():
+	case <-timer.C:
 		return nil, status.Error(codes.DeadlineExceeded, "timed out waiting for VM shutdown")
 	case err = <-shutdownCh:
 		if err != nil {

runtime/service_integ_test.go

@@ -441,7 +441,7 @@ func TestMultipleVMs_Isolated(t *testing.T) {
 
 			containerWg.Wait()
 
-			_, err = fcClient.StopVM(ctx, &proto.StopVMRequest{VMID: strconv.Itoa(vmID), VMShutdownTimeoutSeconds: 10})
+			_, err = fcClient.StopVM(ctx, &proto.StopVMRequest{VMID: strconv.Itoa(vmID), VMShutdownTimeoutSeconds: 5})
 			require.NoError(t, err, "failed to stop VM %d", vmID)
 		}(vmID)
 	}