#!/usr/bin/env python3

# Copyright (c) 2019 Arista Networks, Inc.  All rights reserved.
# Arista Networks, Inc. Confidential and Proprietary.

'''
Dynamically configure some sysctl values based on things like total memory.
'''

import os
import re
import subprocess
import sys

import Ark

KB = 1024
MB = 1024 * KB
GB = 1024 * MB

CPU_DEVICES_DIR = '/sys/devices/system/cpu'
PROC_SYS_BASE_PATH = '/proc/sys'

def _setProcSys( path, value, basePath=PROC_SYS_BASE_PATH ):
   fullPath = os.path.join( basePath, path )
   value = str( value )

   try:
      with open( fullPath, 'w' ) as f:
         f.write( value )
   except OSError as e:
      print( f'Failed to set {fullPath} to {value}: {e}', file=sys.stderr )

def toggleNmiWatchdog():
   # See BUG90059 for an indepth explanation as to why we do this
   if Ark.getPlatform() != 'crow':
      return

   _setProcSys( 'kernel/watchdog', 0 )
   _setProcSys( 'kernel/watchdog', 1 )

def setMinFreeMem( memSize ):
   # Set vm.min_free_kbytes based on amount of memory installed
   if memSize <= 2*GB:
      minFreeKB = 16*KB
   elif memSize <= 4*GB:
      minFreeKB = 32*KB
   elif memSize <= 8*GB:
      minFreeKB = 64*KB
   elif memSize <= 16*GB:
      minFreeKB = 128*KB
   else:
      minFreeKB = 256*KB

   _setProcSys( 'vm/min_free_kbytes', minFreeKB )

def enableEcc( memSize ):
   _setProcSys( 'module/edac_core/parameters/edac_mc_panic_on_ue', 1,
                basePath='/sys' )

   # Ensure ECC demand scrubbing is enabled.
   amdVendor = '1022'
   amdF3List = [ '1403', # PCI_DEVICE_ID_AMD_15H_M10H_F3
                 '141d', # PCI_DEVICE_ID_AMD_15H_M30H_NB_F3
                 '1573', # PCI_DEVICE_ID_AMD_15H_M60H_NB_F3
                 '1603', # PCI_DEVICE_ID_AMD_15H_NB_F3
                 '1533', # PCI_DEVICE_ID_AMD_16H_NB_F3
                 '1583'  # PCI_DEVICE_ID_AMD_16H_M30H_NB_F3
               ]
   configOffset = '0x5c.b'

   for f3Device in amdF3List:
      pciAddresses = subprocess.check_output( [ 'lspci', '-d',
         f'{amdVendor}:{f3Device}' ], text=True )
      for pciAddress in pciAddresses.split( '\n' ):
         if not pciAddress:
            continue
         pciAddress = re.search( r'^[0-9a-fA-F]{2}:[0-9a-fA-F]{2}.[0-9a-fA-F]',
                                 pciAddress ).group( 0 )
         subprocess.check_call( [ 'setpci', '-s', pciAddress,
                                  f'{configOffset}=0x01' ] )

   if not os.path.exists( '/sys/devices/system/edac/mc/mc0/sdram_scrub_rate' ):
      return

   # Configure the kernel and ECC devices. This includes:
   # - Enable ECC patrol scrubbing which will enable ECC demand scrubbing if it is
   #   not already enabled.
   if memSize <= 2*GB:
      eccScrubRate = 24427
   elif memSize <= 4*GB:
      eccScrubRate = 48854
   elif memSize <= 8*GB:
      eccScrubRate = 97650
   elif memSize <= 16*GB:
      eccScrubRate = 195300
   else:
      eccScrubRate = 390720

   _setProcSys( 'devices/system/edac/mc/mc0/sdram_scrub_rate', eccScrubRate,
                basePath='/sys' )

def memBasedConfig():
   # Configure the kernel and ECC devices based on amount of memory installed
   try:
      with open( '/proc/meminfo' ) as f:
         for line in f:
            if 'MemTotal' in line:
               memSize = int( line.split() [ 1 ] ) * KB
               break
   except OSError:
      memSize = 8*GB

   setMinFreeMem( memSize )
   enableEcc( memSize )

def cpuFreqConfig():
   ''' Make sure all the CPUs are configured with the performance governor '''

   if not os.path.exists( '/sys/devices/system/cpu/intel_pstate' ):
      return

   for cpuDir in os.listdir( CPU_DEVICES_DIR ):
      if not cpuDir.startswith( 'cpu' ):
         continue

      governorPath = os.path.join( CPU_DEVICES_DIR, cpuDir,
                                   'cpufreq/scaling_governor' )
      if os.path.exists( governorPath ):
         _setProcSys( governorPath, 'performance', basePath='/' )

if __name__ == "__main__":
   toggleNmiWatchdog()
   memBasedConfig()
   cpuFreqConfig()
