#!/usr/bin/env python

# 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.
'''

from __future__ import print_function

import os
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 ( IOError, OSError ) as e:
      print( 'Failed to set %s to %s: %s' % ( fullPath, 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 configureRssCanary( memSize ):
   if memSize <= 2*GB:
      rssSize = 1024
   elif memSize <= 4*GB:
      rssSize = 1500
   elif memSize <= 8*GB:
      rssSize = 2500
   else:
      rssSize = 3500

   _setProcSys( 'kernel/rss_warn_thresh', rssSize )

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

   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 )
   configureRssCanary( 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 ):
         with open( governorPath, 'w' ) as f:
            f.write( 'performance' )

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